Wire witness support inside Megatron

.gitignore

@@ -16,6 +16,7 @@ onelogger.err
 runs/
 /test_cases/
 **/dist/
+.idea
 
 # Sphinx documentation
 docs/_build

megatron/core/models/gpt/gpt_model.py

@@ -553,6 +553,7 @@ def forward(
         loss_mask: Optional[Tensor] = None,
         padding_mask: Optional[Tensor] = None,
         mtp_kwargs: Optional[dict] = {},
+        witness_ids: Optional[Tensor] = None,
     ) -> Tensor:
         """Forward function of the GPT Model This function passes the input tensors
         through the embedding layer, and then the decoder and finally into the post
@@ -592,6 +593,12 @@ def forward(
 
         rotary_pos_cos_sin = preproc_output[6] if len(preproc_output) == 7 else None
 
+        if witness_ids is not None and hasattr(self, "local_head_witness"):
+            if decoder_input is not None:
+                decoder_input = self.local_head_witness(witness_ids, decoder_input)
+            else:
+                self.decoder.input_tensor = self.local_head_witness(witness_ids, self.decoder.input_tensor)
+
         # Run decoder.
         hidden_states = self.decoder(
             hidden_states=decoder_input,
@@ -608,6 +615,9 @@ def forward(
             **(extra_block_kwargs or {}),
         )
 
+        if witness_ids is not None and hasattr(self, "local_tail_witness"):
+            hidden_states = self.local_tail_witness(witness_ids, hidden_states)
+
         return self._postprocess(
             hidden_states=hidden_states,
             input_ids=input_ids,

megatron/core/optimizer/distrib_optimizer.py

@@ -407,8 +407,14 @@ def _build_model_and_main_param_groups(
                         )
                         if hasattr(model_param, 'shared'):
                             shard_main_param.shared = model_param.shared
+                        if getattr(model_param, '_is_witness_param', False):
+                            shard_main_param._is_witness_param = True
                     else:
                         # When using precision-aware optimizer, main params are held by FusedAdam.
+                        assert not getattr(model_param, '_is_witness_param', False), (
+                            "Witness params are not supported with precision-aware optimizer. "
+                            "_is_witness_param flag cannot be propagated to FusedAdam-held params."
+                        )
                         shard_main_param = None
 
                     # Store handle to main_param.
@@ -431,6 +437,8 @@ def _build_model_and_main_param_groups(
                     )
                     if hasattr(model_param, 'shared'):
                         shard_model_param.shared = model_param.shared
+                    if getattr(model_param, '_is_witness_param', False):
+                        shard_model_param._is_witness_param = True
 
                 else:
                     raise TypeError(

megatron/core/optimizer/optimizer.py

@@ -156,7 +156,8 @@ def get_main_grads_for_grad_norm(self) -> List[torch.Tensor]:
             is_not_tp_duplicate = tensor_parallel.param_is_not_tensor_parallel_duplicate(
                 param, getattr(self, 'tp_group', None)
             )
-            if grad_not_none and is_not_shared and is_not_tp_duplicate:
+            is_not_witness = not getattr(param, "_is_witness_param", False)
+            if grad_not_none and is_not_shared and is_not_tp_duplicate and is_not_witness:
                 grads_for_norm.append(grad)
 
         return grads_for_norm