changed mup_m_width to mup_width_multiplier

configs/neox_arguments.md

@@ -1567,7 +1567,7 @@ Training Arguments
 
 
 
-- **mup_m_width**: int
+- **mup_width_multiplier**: int
 
  Default = 1
 
@@ -1577,7 +1577,7 @@ Training Arguments
 
 - **mup_d_model_base**: int
 
- Default = 64
+ Default = 256
 
  d_model,base
  Proxy (base) model's layer width

megatron/learning_rates.py

@@ -37,7 +37,7 @@ def __init__(
  use_checkpoint_lr_scheduler=True,
  override_lr_scheduler=False,
  use_mup=False,
- mup_m_width=1,
+ mup_width_multiplier=1,
  ):
 
  # Class values.
@@ -52,7 +52,7 @@ def __init__(
  self.override_lr_scheduler = override_lr_scheduler
  self.use_checkpoint_lr_scheduler = use_checkpoint_lr_scheduler
  self.use_mup = use_mup
- self.mup_m_width = mup_m_width
+ self.mup_width_multiplier = mup_width_multiplier
  if self.override_lr_scheduler:
  assert not self.use_checkpoint_lr_scheduler, (
  "both override and " "use-checkpoint are set."
@@ -98,7 +98,7 @@ def step(self, step_num=None):
  new_lr = self.get_lr()
  for group in self.optimizer.param_groups:
  if self.use_mup and ("lr_adjust" in group) and group["lr_adjust"] is True:
- group["lr"] = new_lr / self.mup_m_width
+ group["lr"] = new_lr / self.mup_width_multiplier
  else:
  group["lr"] = new_lr
 

megatron/model/init_functions.py

@@ -68,12 +68,12 @@ def _orthogonal(tensor, gain=1):
  return tensor
 
 
-def orthogonal_init_method(n_layers=1, mup_m_width=1.0):
+def orthogonal_init_method(n_layers=1, mup_width_multiplier=1.0):
  """Fills the input Tensor with a (semi) orthogonal matrix, as described in
  Exact solutions to the nonlinear dynamics of learning in deep linear neural networks - Saxe, A. et al. (2013)
  Optionally scaling by number of layers possible, as introduced in OBST - Nestler et. al. (2021, to be released)"""
 
- if mup_m_width != 1:
+ if mup_width_multiplier != 1:
  raise ValueError(
  "Orthogonal init needs to be patched to support mup. Disable mup or use a different init method to avoid this error"
  )
@@ -84,57 +84,57 @@ def init_(tensor):
  return init_
 
 
-def xavier_uniform_init_method(mup_m_width=1.0):
+def xavier_uniform_init_method(mup_width_multiplier=1.0):
  """Fills the input Tensor with values according to the method described in Understanding the difficulty of
  training deep feedforward neural networks - Glorot, X. & Bengio, Y. (2010), using a uniform distribution."""
 
- def init_(tensor, mup_m_width=mup_m_width):
+ def init_(tensor, mup_width_multiplier=mup_width_multiplier):
  init_weight = torch.nn.init.xavier_uniform_(tensor)
- if mup_m_width != 1:
+ if mup_width_multiplier != 1:
  with torch.no_grad():
- init_weight.div_(mup_m_width)
+ init_weight.div_(mup_width_multiplier)
  return init_weight
 
  return init_
 
 
-def xavier_normal_init_method(mup_m_width=1.0):
+def xavier_normal_init_method(mup_width_multiplier=1.0):
  """Fills the input Tensor with values according to the method described in Understanding the difficulty of
  training deep feedforward neural networks - Glorot, X. & Bengio, Y. (2010), using a normal distribution."""
 
- def init_(tensor, mup_m_width=mup_m_width):
+ def init_(tensor, mup_width_multiplier=mup_width_multiplier):
  init_weight = torch.nn.init.xavier_normal_(tensor)
- if mup_m_width != 1:
+ if mup_width_multiplier != 1:
  with torch.no_grad():
- init_weight.div_(mup_m_width)
+ init_weight.div_(mup_width_multiplier)
  return init_weight
 
  return init_
 
 
-def small_init_init_method(dim, mup_m_width=1.0):
+def small_init_init_method(dim, mup_width_multiplier=1.0):
  """Fills the input Tensor with values according to the method described in Transformers without Tears: Improving
  the Normalization of Self-Attention - Nguyen, T. & Salazar, J. (2010), using a normal distribution."""
  std = math.sqrt(2 / (5 * dim))
 
- def init_(tensor, mup_m_width=mup_m_width):
+ def init_(tensor, mup_width_multiplier=mup_width_multiplier):
  init_weight = torch.nn.init.normal_(tensor, mean=0.0, std=std)
- if mup_m_width != 1:
+ if mup_width_multiplier != 1:
  with torch.no_grad():
- init_weight.div_(mup_m_width)
+ init_weight.div_(mup_width_multiplier)
  return init_weight
 
  return init_
 
 
-def wang_init_method(n_layers, dim, mup_m_width=1.0):
+def wang_init_method(n_layers, dim, mup_width_multiplier=1.0):
  std = 2 / n_layers / math.sqrt(dim)
 
- def init_(tensor, mup_m_width=mup_m_width):
+ def init_(tensor, mup_width_multiplier=mup_width_multiplier):
  init_weight = torch.nn.init.normal_(tensor, mean=0.0, std=std)
- if mup_m_width != 1:
+ if mup_width_multiplier != 1:
  with torch.no_grad():
- init_weight.div_(mup_m_width)
+ init_weight.div_(mup_width_multiplier)
  return init_weight
 
  return init_
@@ -145,30 +145,30 @@ def get_init_methods(args):
  def _get(name):
  if name == "normal":
  return init_method_normal(
- sigma=args.init_method_std/math.sqrt(args.mup_m_width)
+ sigma=args.init_method_std/math.sqrt(args.mup_width_multiplier)
  )
  elif name == "scaled_normal":
  return scaled_init_method_normal(
- sigma=args.init_method_std/math.sqrt(args.mup_m_width),
+ sigma=args.init_method_std/math.sqrt(args.mup_width_multiplier),
  num_layers=args.num_layers
  )
  elif name == "orthogonal":
- return orthogonal_init_method(args.mup_m_width)
+ return orthogonal_init_method(args.mup_width_multiplier)
  elif name == "scaled_orthogonal":
  return orthogonal_init_method(
- args.num_layers, args.mup_m_width
+ args.num_layers, args.mup_width_multiplier
  )
  elif name == "xavier_uniform":
- return xavier_uniform_init_method(args.mup_m_width)
+ return xavier_uniform_init_method(args.mup_width_multiplier)
  elif name == "xavier_normal":
- return xavier_normal_init_method(args.mup_m_width)
+ return xavier_normal_init_method(args.mup_width_multiplier)
  elif name == "wang_init":
  return wang_init_method(
- args.num_layers, args.hidden_size, args.mup_m_width
+ args.num_layers, args.hidden_size, args.mup_width_multiplier
  )
  elif name == "small_init":
  return small_init_init_method(
- args.hidden_size, args.mup_m_width
+ args.hidden_size, args.mup_width_multiplier
  )
  else:
  raise NotImplementedError(f"Unknown init method {name}")

megatron/model/transformer.py

@@ -972,7 +972,7 @@ def parallel_lm_logits(input_, word_embeddings_weight, parallel_output, bias=Non
  logits_parallel = F.linear(input_parallel, word_embeddings_weight, bias)
 
  if args is not None and args.use_mup:
- logits_parallel /= args.mup_m_width
+ logits_parallel /= args.mup_width_multiplier
 
  # Gather if needed.
  if parallel_output:

megatron/mpu/layers.py

@@ -429,7 +429,7 @@ def __init__(
  self.stride = stride
  self.mup_rescale_parameters = mup_rescale_parameters
  self.use_mup = neox_args.use_mup
- self.m_width = neox_args.mup_m_width
+ self.m_width = neox_args.mup_width_multiplier
 
  # Parameters.
  # Note: torch.nn.functional.linear performs XA^T + b and as a result
@@ -627,7 +627,7 @@ def __init__(
  self.keep_master_weight_for_test = keep_master_weight_for_test
  self.mup_rescale_parameters = mup_rescale_parameters
  self.use_mup = neox_args.use_mup
- self.m_width = neox_args.mup_m_width
+ self.m_width = neox_args.mup_width_multiplier
 
  # Parameters.
  # Note: torch.nn.functional.linear performs XA^T + b and as a result

megatron/neox_arguments/neox_args.py

@@ -1046,7 +1046,7 @@ class NeoXArgsTraining(NeoXArgsTemplate):
  Embedding output multiplier
  """
 
- mup_m_width: float = None
+ mup_width_multiplier: float = None
  """
  Manually set the layer width multiplier (d_model/d_model,base)
  """

megatron/training.py

@@ -422,9 +422,9 @@ def get_model(neox_args, use_cache=False):
  # neox_args.use_mup = False
  if neox_args.use_mup:
 
- if neox_args.mup_m_width == 1:
- neox_args.mup_m_width = neox_args.hidden_size / neox_args.mup_d_model_base
- print_rank_0(f"mup_m_width set to {neox_args.mup_m_width}")
+ if neox_args.mup_width_multiplier == 1:
+ neox_args.mup_width_multiplier = neox_args.hidden_size / neox_args.mup_d_model_base
+ print_rank_0(f"mup_width_multiplier set to {neox_args.mup_width_multiplier}")
 
  # base_shapes = f"{neox_args.base_shapes_file}.{torch.distributed.get_rank()}"
 
@@ -640,7 +640,7 @@ def get_learning_rate_scheduler(optimizer, neox_args):
  use_checkpoint_lr_scheduler=neox_args.use_checkpoint_lr_scheduler,
  override_lr_scheduler=neox_args.override_lr_scheduler,
  use_mup=neox_args.use_mup,
- mup_m_width=neox_args.mup_m_width,
+ mup_width_multiplier=neox_args.mup_width_multiplier,
  )
 
  return lr_scheduler