Alias core.parallel_state as mpu and use it throughout code. RIP mpu.

megatron/checkpointing.py

@@ -9,8 +9,8 @@
 
 import torch
 
-from megatron import (core,
- update_num_microbatches)
+from megatron import update_num_microbatches
+from megatron.core import mpu, tensor_parallel
 from .global_vars import get_args
 from .utils import (unwrap_model,
  print_rank_0)
@@ -79,11 +79,11 @@ def get_checkpoint_names(checkpoints_path, iteration, use_distributed_optimizer,
 
  # Use both the tensor and pipeline MP rank.
  if pipeline_parallel is None:
- pipeline_parallel = (core.get_pipeline_model_parallel_world_size() > 1)
+ pipeline_parallel = (mpu.get_pipeline_model_parallel_world_size() > 1)
  if tensor_rank is None:
- tensor_rank = core.get_tensor_model_parallel_rank()
+ tensor_rank = mpu.get_tensor_model_parallel_rank()
  if pipeline_rank is None:
- pipeline_rank = core.get_pipeline_model_parallel_rank()
+ pipeline_rank = mpu.get_pipeline_model_parallel_rank()
 
  # Use both the tensor and pipeline MP rank. If using the distributed
  # optimizer, then the optimizer's path must additionally include the
@@ -98,7 +98,7 @@ def get_checkpoint_names(checkpoints_path, iteration, use_distributed_optimizer,
  if use_distributed_optimizer:
  model_name = os.path.join(common_path, "model_rng.pt")
  optim_name = os.path.join(
- common_path + "_%03d" % core.get_data_parallel_rank(),
+ common_path + "_%03d" % mpu.get_data_parallel_rank(),
  "optim.pt")
  else:
  model_name = optim_name = os.path.join(common_path, "model_optim_rng.pt")
@@ -185,18 +185,18 @@ def get_rng_state():
  'np_rng_state': np.random.get_state(),
  'torch_rng_state': torch.get_rng_state(),
  'cuda_rng_state': torch.cuda.get_rng_state(),
- 'rng_tracker_states': core.tensor_parallel.get_cuda_rng_tracker().get_states()}
+ 'rng_tracker_states': tensor_parallel.get_cuda_rng_tracker().get_states()}
 
  rng_state_list = None
  if torch.distributed.is_initialized() and \
- core.get_data_parallel_world_size() > 1 and \
+ mpu.get_data_parallel_world_size() > 1 and \
  args.data_parallel_random_init:
  rng_state_list = \
- [None for i in range(core.get_data_parallel_world_size())]
+ [None for i in range(mpu.get_data_parallel_world_size())]
  torch.distributed.all_gather_object(
  rng_state_list,
  rng_state,
- group=core.get_data_parallel_group())
+ group=mpu.get_data_parallel_group())
  else:
  rng_state_list = [rng_state]
 
@@ -223,7 +223,7 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler):
  # Collect args, model, RNG.
  model_state_dict = {}
  if not torch.distributed.is_initialized() \
- or core.get_data_parallel_rank() == 0:
+ or mpu.get_data_parallel_rank() == 0:
 
  # Arguments, iteration, and model.
  model_state_dict['args'] = args
@@ -233,7 +233,7 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler):
  model_state_dict['model'] = model[0].state_dict_for_save_checkpoint()
  else:
  for i in range(len(model)):
- core.set_virtual_pipeline_model_parallel_rank(i)
+ mpu.set_virtual_pipeline_model_parallel_rank(i)
  model_state_dict['model%d' % i] = \
  model[i].state_dict_for_save_checkpoint()
 
@@ -246,7 +246,7 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler):
  optim_state_dict = {}
  if not args.no_save_optim \
  and (not torch.distributed.is_initialized()
- or core.get_data_parallel_rank() == 0
+ or mpu.get_data_parallel_rank() == 0
  or args.use_distributed_optimizer):
 
  # Optimizer stuff.
@@ -548,7 +548,7 @@ def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', stri
  model[0].load_state_dict(model_state_dict['model'], strict=strict)
  else:
  for i in range(len(model)):
- core.set_virtual_pipeline_model_parallel_rank(i)
+ mpu.set_virtual_pipeline_model_parallel_rank(i)
  model[i].load_state_dict(model_state_dict['model%d' % i], strict=strict)
 
  # Fix up query/key/value matrix ordering if needed
@@ -580,7 +580,7 @@ def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', stri
  # access rng_state for data parallel rank
  if args.data_parallel_random_init:
 
- rng_state = model_state_dict['rng_state'][core.get_data_parallel_rank()]
+ rng_state = model_state_dict['rng_state'][mpu.get_data_parallel_rank()]
  else:
  rng_state = model_state_dict['rng_state'][0]
  random.setstate(rng_state['random_rng_state'])
@@ -590,7 +590,7 @@ def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', stri
  # Check for empty states array
  if not rng_state['rng_tracker_states']:
  raise KeyError
- core.tensor_parallel.get_cuda_rng_tracker().set_states(
+ tensor_parallel.get_cuda_rng_tracker().set_states(
  rng_state['rng_tracker_states'])
  else: # backward compatability
  random.setstate(model_state_dict['random_rng_state'])
@@ -600,7 +600,7 @@ def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', stri
  # Check for empty states array
  if not model_state_dict['rng_tracker_states']:
  raise KeyError
- core.tensor_parallel.get_cuda_rng_tracker().set_states(
+ tensor_parallel.get_cuda_rng_tracker().set_states(
  model_state_dict['rng_tracker_states'])
  except KeyError:
  print_rank_0('Unable to load rng state from checkpoint {}. '
@@ -640,7 +640,7 @@ def load_biencoder_checkpoint(model, only_query_model=False,
  args.use_distributed_optimizer,
  release=False)
 
- if core.get_data_parallel_rank() == 0:
+ if mpu.get_data_parallel_rank() == 0:
  print('global rank {} is loading checkpoint {}'.format(
  torch.distributed.get_rank(), checkpoint_name))
 
@@ -656,7 +656,7 @@ def load_biencoder_checkpoint(model, only_query_model=False,
  model[0].load_state_dict(ret_state_dict)
  torch.distributed.barrier()
 
- if core.get_data_parallel_rank() == 0:
+ if mpu.get_data_parallel_rank() == 0:
  print(' successfully loaded {}'.format(checkpoint_name))
 
  return model

megatron/core/__init__.py

@@ -1,13 +1,12 @@
-from .parallel_state import (
- initialize_model_parallel,
- get_tensor_model_parallel_world_size,
- get_tensor_model_parallel_rank,
- get_pipeline_model_parallel_world_size,
- get_pipeline_model_parallel_rank,
- get_virtual_pipeline_model_parallel_rank, set_virtual_pipeline_model_parallel_rank,
- get_data_parallel_world_size,
- get_data_parallel_rank,
- get_global_memory_buffer,
- get_num_layers,
-)
-from megatron.core import tensor_parallel
+import megatron.core.parallel_state
+import megatron.core.tensor_parallel
+import megatron.core.utils
+
+# Alias parallel_state as mpu, its legacy name
+mpu = parallel_state
+
+__all__ = [
+ "parallel_state",
+ "tensor_parallel",
+ "utils",
+]

megatron/core/parallel_state.py

@@ -47,10 +47,6 @@
 # Memory buffers to avoid dynamic memory allocation
 _GLOBAL_MEMORY_BUFFER = None
 
-def is_unitialized():
- """Useful for code segments that may be accessed with or without mpu initialization"""
- return _DATA_PARALLEL_GROUP is None
-
 
 def initialize_model_parallel(
  tensor_model_parallel_size: int = 1,

megatron/core/tensor_parallel/__init__.py

@@ -5,6 +5,7 @@
  ColumnParallelLinear,
  RowParallelLinear,
  VocabParallelEmbedding,
+ set_tensor_model_parallel_attributes,
  set_defaults_if_not_set_tensor_model_parallel_attributes,
  copy_tensor_model_parallel_attributes,
  param_is_not_tensor_parallel_duplicate,
@@ -23,10 +24,14 @@
 from .random import (
  checkpoint,
  get_cuda_rng_tracker,
- model_parallel_cuda_manual_seed
+ model_parallel_cuda_manual_seed,
 )
 
-from .utils import split_tensor_along_last_dim
+from .utils import (
+ split_tensor_along_last_dim,
+ split_tensor_into_1d_equal_chunks,
+ gather_split_1d_tensor,
+)
 
 __all__ = [
  # cross_entropy.py

megatron/core/tensor_parallel/random.py

@@ -17,6 +17,10 @@
  get_tensor_model_parallel_world_size,
 )
 
+from .utils import (
+ split_tensor_into_1d_equal_chunks,
+ gather_split_1d_tensor,
+)
 
 # Default name for the model parallel rng tracker.
 _MODEL_PARALLEL_RNG_TRACKER_NAME = 'model-parallel-rng'
@@ -55,38 +59,6 @@ def cb():
  _lazy_call(cb)
 
 
-def split_tensor_into_1d_equal_chunks(tensor, new_buffer=False):
- """Break a tensor into equal 1D chunks."""
- partition_size = torch.numel(tensor) // \
- get_tensor_model_parallel_world_size()
- start_index = partition_size * get_tensor_model_parallel_rank()
- end_index = start_index + partition_size
- if new_buffer:
- data = torch.empty(partition_size, dtype=tensor.dtype,
- device=torch.cuda.current_device(),
- requires_grad=False)
- data.copy_(tensor.view(-1)[start_index:end_index])
- else:
- data = tensor.view(-1)[start_index:end_index]
- return data
-
-
-def gather_split_1d_tensor(tensor):
- """Opposite of above function, gather values from model parallel ranks."""
- numel_gathered = torch.numel(tensor) * \
- get_tensor_model_parallel_world_size()
- gathered = torch.empty(numel_gathered, dtype=tensor.dtype,
- device=torch.cuda.current_device(),
- requires_grad=False)
- # TODO: This API is experimental in pytorch (as of Feb 2022) and
- # this might break in future pytorch releases. We chose this API
- # as opposed to torch.distributed.all_gather for efficiency reasons.
- # This API calls directly NCCL all-gather versus the former does
- # internal copies and can potentially cause slow down.
- torch.distributed._all_gather_base(gathered, tensor,
- group=get_tensor_model_parallel_group())
- return gathered
-
 
 class CudaRNGStatesTracker:
  """Tracker for the cuda RNG states.

megatron/core/tensor_parallel/utils.py

@@ -10,12 +10,16 @@ def split_tensor_along_last_dim(
  num_partitions: int,
  contiguous_split_chunks: bool = False,
 ) -> List[torch.Tensor]:
- """Split a tensor along its last dimension.
- Arguments:
- tensor: input tensor.
- num_partitions: number of partitions to split the tensor
- contiguous_split_chunks: If True, make each chunk contiguous
- in memory.
+ """ Split a tensor along its last dimension.
+
+ Arguments:
+ tensor: input tensor.
+ num_partitions: number of partitions to split the tensor
+ contiguous_split_chunks: If True, make each chunk contiguous
+ in memory.
+
+ Returns:
+ A list of Tensors
  """
  # Get the size and dimension.
  last_dim = tensor.dim() - 1
@@ -28,11 +32,64 @@ def split_tensor_along_last_dim(
 
  return tensor_list
 
+def split_tensor_into_1d_equal_chunks(tensor, new_buffer=False):
+ """ Break a tensor into equal 1D chunks across tensor parallel ranks.
+
+ Returns a Tensor or View with this rank's portion of the data.
+
+ Arguments:
+ tensor: The tensor to split
+
+ Keyword Arguments:
+ new_buffer (bool): If True, returns a new Tensor.
+ If False, returns a view into the existing Tensor.
+ Default is False
+
+ """
+ partition_size = torch.numel(tensor) // \
+ get_tensor_model_parallel_world_size()
+ start_index = partition_size * get_tensor_model_parallel_rank()
+ end_index = start_index + partition_size
+ if new_buffer:
+ data = torch.empty(partition_size, dtype=tensor.dtype,
+ device=torch.cuda.current_device(),
+ requires_grad=False)
+ data.copy_(tensor.view(-1)[start_index:end_index])
+ else:
+ data = tensor.view(-1)[start_index:end_index]
+ return data
+
+
+def gather_split_1d_tensor(tensor):
+ """ Opposite of split_tensor_into_1d_equal_chunks. Gather values from tensor
+ model parallel ranks.
+
+ Returns a new Tensor with the gathered data.
+
+ Arguments:
+ tensor: A Tensor or view of this rank's portion of the data.
+ """
+ numel_gathered = torch.numel(tensor) * \
+ get_tensor_model_parallel_world_size()
+ gathered = torch.empty(numel_gathered, dtype=tensor.dtype,
+ device=torch.cuda.current_device(),
+ requires_grad=False)
+ # TODO: This API is experimental in pytorch (as of Feb 2022) and
+ # this might break in future pytorch releases. We chose this API
+ # as opposed to torch.distributed.all_gather for efficiency reasons.
+ # This API calls directly NCCL all-gather versus the former does
+ # internal copies and can potentially cause slow down.
+ torch.distributed._all_gather_base(gathered, tensor,
+ group=get_tensor_model_parallel_group())
+ return gathered
+
 
 class VocabUtility:
- """Split the vocabulary into `world_size` chunks and return the
- first and last index of the vocabulary belonging to the `rank`
- partition: Note that indices in [fist, last)"""
+ """ Split the vocabulary into `world_size` chunks and return the first
+ and last index of the vocabulary belonging to the `rank`
+ partition: Note that indices in [fist, last)
+
+ """
 
  @staticmethod
  def vocab_range_from_per_partition_vocab_size(

megatron/core/utils.py

@@ -21,35 +21,6 @@ def divide(numerator, denominator):
  return numerator // denominator
 
 
-def split_tensor_into_1d_equal_chunks(tensor):
- """Break a tensor into equal 1D chunks."""
- data = tensor.view(-1)
- partition_size = (
- torch.numel(data) // parallel_state.get_tensor_model_parallel_world_size()
- )
- start_index = partition_size * parallel_state.get_tensor_model_parallel_rank()
- end_index = start_index + partition_size
- return data[start_index:end_index]
-
-
-def gather_split_1d_tensor(tensor):
- """Opposite of above function, gather values from model parallel ranks."""
- world_size = parallel_state.get_tensor_model_parallel_world_size()
- numel = torch.numel(tensor)
- numel_gathered = world_size * numel
- gathered = torch.empty(
- numel_gathered,
- dtype=tensor.dtype,
- device=torch.cuda.current_device(),
- requires_grad=False,
- )
- torch.distributed._all_gather_base(
- gathered,
- tensor,
- group=parallel_state.get_tensor_model_parallel_group()
- )
- return gathered
-
 class GlobalMemoryBuffer:
  """Global buffer to avoid dynamic memory allocations.
  Caller should ensure that buffers of the same name

megatron/data/biencoder_dataset_utils.py

@@ -4,7 +4,8 @@
 import numpy as np
 import torch
 
-from megatron import get_args, get_tokenizer, mpu, print_rank_0
+from megatron import get_args, get_tokenizer, print_rank_0
+from megatron.core import mpu, tensor_parallel
 from megatron.data.dataset_utils import create_masked_lm_predictions, \
  pad_and_convert_to_numpy
 from megatron.data.data_samplers import MegatronPretrainingSampler
@@ -57,7 +58,7 @@ def get_ict_batch(data_iterator):
  data = None
  else:
  data = next(data_iterator)
- data_b = mpu.broadcast_data(keys, data, datatype)
+ data_b = tensor_parallel.broadcast_data(keys, data, datatype)
 
  # Unpack.
  query_tokens = data_b['query_tokens'].long()