docstr: Train index (rank 0's method).
docstr_tok: ["Train","index","(","rank","0","'s","method",")","."]
code: def _train(self):
      """Train index (rank 0's method)."""
  
      args = get_retro_args()
  
      assert torch.distributed.get_rank() == 0
  
      # Set num threads (torch.distributed reset it to 1).
      # faiss.omp_set_num_threads(32)
      faiss.omp_set_num_threads(64)
      # faiss.omp_set_num_threads(128)
  
      empty_index_path = self.get_empty_index_path()
  
      # Index already exists? -> return.
      if os.path.isfile(empty_index_path):
          return
  
      # Load data.
      merged_path = get_training_data_merged_path()
      inp = np.memmap(
          merged_path,
          dtype="f4",
          mode="r",
      ).reshape((-1, args.hidden_size))
  
      # Init index.
      index = faiss.index_factory(
          args.retro_index_nfeats, args.retro_index_str
      )
  
      # Move to GPU.
      print("> move faiss index to gpu.")
      index_ivf = faiss.extract_index_ivf(index)
      clustering_index = faiss.index_cpu_to_all_gpus(
          faiss.IndexFlatL2(index_ivf.d)
      )
      index_ivf.clustering_index = clustering_index
      print("> finished moving to gpu.")
      self.c_verbose(index, True)
      self.c_verbose(index_ivf, True)
      self.c_verbose(index_ivf.quantizer, True)
      self.c_verbose(index_ivf.clustering_index, True)
  
      # Train index.
      index.train(inp)
  
      # Save index.
      faiss.write_index(index, empty_index_path)
  
file_path: megatron-lm/tools/retro/index/indexes/faiss_base.py
imports: from datetime import timedelta
  import numpy
  import os
  import torch
  from tqdm import tqdm
  from megatron import get_retro_args, print_rank_0
  from tools.bert_embedding import BertEmbedder
  from tools.retro.external_libs import faiss
  from tools.retro.index.index import Index
  from tools.retro.index.utils import get_training_data_merged_path, num_samples_to_block_ranges
label_desc: Use image node_id 1 for calling the FaissBaseIndex obj's underlying member method code with example usage: obj._train() without return types