Improve our mup implementation

configs/neox_arguments.md

@@ -111,7 +111,7 @@ Logging Arguments
 
 - **git_hash**: str
 
-    Default = ebaeec1
+    Default = 5f09348
 
     current git hash of repository
 
@@ -1548,6 +1548,14 @@ Training Arguments
 
 
 
+- **mup_deferred_init**: bool
+
+    Default = False
+
+    Whether to fully initialize the base and delta models (set to true for big target models)
+
+
+
 ## NeoXArgsDeepspeedConfig
 
 Args for deepspeed config

megatron/model/gpt2_model.py

@@ -222,6 +222,9 @@ def init_specs(self):
                 heads=self.neox_args.num_attention_heads,
             )
 
+        if self.neox_args.use_mup and self.neox_args.mup_input_temp is not None:
+            self.specs.append(lambda x: x * self.neox_args.mup_input_temp)
+
         # Transformer layers
         for i in range(self.neox_args.num_layers):
             layer_type = self.neox_args.attention_config[i]
@@ -260,6 +263,9 @@ def init_specs(self):
             LayerSpec(NormPipe, norm, self.neox_args.hidden_size, eps=eps)
         )
 
+        if self.neox_args.use_mup and self.neox_args.output_temp is not None:
+            self.specs.append(lambda x: x * self.neox_args.mup_output_temp / self.neox_args.hidden_size)
+
         # outputs are now a single tensor: hidden_states
 
         def _logits_helper(embedding, lm_output):

megatron/model/transformer.py

@@ -227,7 +227,10 @@ def __init__(
         )
 
         coeff = None
-        self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
+        if neox_args.use_mup:
+            self.norm_factor = self.hidden_size_per_attention_head / neox_args.mup_attn_temp
+        else:
+            self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
         if self.apply_query_key_layer_scaling:
             coeff = max(1, self.layer_number)
             self.norm_factor *= coeff

megatron/neox_arguments/neox_args.py

@@ -1030,6 +1030,11 @@ class NeoXArgsTraining(NeoXArgsTemplate):
     What to scale width by when creating the delta model for mup
     """
 
+    mup_deferred_init: bool = False
+    """
+    Whether to fully initialize the base and delta models (set to true for big target models)
+    """
+
 
 @dataclass
 class NeoXArgsTextgen(NeoXArgsTemplate):

megatron/training.py

@@ -76,13 +76,35 @@ def save_base_shapes(neox_args, base_shapes, use_cache):
     # Instantiation of the base model fails in the init function (init_functions.py) because we haven't called set_base_shapes on it at this point, so disable it temporarily here
     neox_args.use_mup = False
 
-    base_model = GPT2ModelPipe(
-        neox_args=neox_args,
-        num_tokentypes=0,
-        parallel_output=True,
-        topology=mpu.get_topology(),
-        use_cache=use_cache,
-    )
+    # print(
+    #     f"MEM BEFORE BASE MUP: {torch.cuda.memory_allocated(device_index)} on rank {torch.distributed.get_rank()}"
+    # )
+    if neox_args.mup_deferred_init:
+        try:
+            from torchdistx import deferred_init
+        except ModuleNotFoundError:
+            print("Please install torchdistx https://github.com/pytorch/torchdistx")
+            raise Exception
+        base_model = deferred_init.deferred_init(
+            GPT2ModelPipe,
+            neox_args=neox_args,
+            num_tokentypes=0,
+            parallel_output=True,
+            topology=mpu.get_topology(),
+            use_cache=use_cache,
+        )
+    else:
+        base_model = GPT2ModelPipe(
+            neox_args=neox_args,
+            num_tokentypes=0,
+            parallel_output=True,
+            topology=mpu.get_topology(),
+            use_cache=use_cache,
+        )
+
+    # print(
+    #     f"MEM AFTER BASE MUP: {torch.cuda.memory_allocated(device_index)} on rank {torch.distributed.get_rank()}"
+    # )
 
     if not neox_args.is_pipe_parallel:
         base_model = base_model.to_sequential()
@@ -100,13 +122,36 @@ def save_base_shapes(neox_args, base_shapes, use_cache):
     old_hidden_size = neox_args.hidden_size
     neox_args.hidden_size = neox_args.hidden_size * neox_args.mup_width_scale
 
-    delta_model = GPT2ModelPipe(
-        neox_args=neox_args,
-        num_tokentypes=0,
-        parallel_output=True,
-        topology=mpu.get_topology(),
-        use_cache=use_cache,
-    )
+    # print(
+    #     f"MEM BEFORE DELTA MUP: {torch.cuda.memory_allocated(device_index)} on rank {torch.distributed.get_rank()}"
+    # )
+    if neox_args.mup_deferred_init:
+        print("Using MUP deferred init")
+        try:
+            from torchdistx import deferred_init
+        except ModuleNotFoundError:
+            print("Please install torchdistx https://github.com/pytorch/torchdistx")
+            raise Exception
+        delta_model = deferred_init.deferred_init(
+            GPT2ModelPipe,
+            neox_args=neox_args,
+            num_tokentypes=0,
+            parallel_output=True,
+            topology=mpu.get_topology(),
+            use_cache=use_cache,
+        )
+    else:
+        delta_model = GPT2ModelPipe(
+            neox_args=neox_args,
+            num_tokentypes=0,
+            parallel_output=True,
+            topology=mpu.get_topology(),
+            use_cache=use_cache,
+        )
+
+    # print(
+    #     f"MEM AFTER BASE MUP: {torch.cuda.memory_allocated(device_index)} on rank {torch.distributed.get_rank()}"
+    # )
 
     if not neox_args.is_pipe_parallel:
         delta_model = delta_model.to_sequential()

requirements/requirements-mup.txt

@@ -0,0 +1 @@
+mup==1.0.0