Support DeepSpeed ZeRO-3 in KDTrainer; fix Liger hidden-states dtype

- Add fully-frozen-model fallback in ModelOptHFTrainer._prepare_model so
  DS ZeRO-3 can prepare a frozen teacher without hitting the empty
  trainable_param_groups assertion.
- Add KDTrainer._ds_gather context manager for explicit param gather,
  since the teacher is loaded under zero.Init but not wrapped in a
  DeepSpeedEngine (no per-module hooks).
- Unify KD sharded Liger compute: delegate student lm_head gather to
  the parent's _sharded_liger_compute and add teacher lm_head gather
  via _apply_teacher_gather.
- Cast outputs.logits to lm_head.weight dtype before Liger fused kernels
  (final RMSNorm may leave hidden_states in fp32).
- Drop redundant KDTrainer._get_lm_head override (inherited).

Signed-off-by: realAsma <akuriparambi@nvidia.com>

Author: realAsma

modelopt/torch/distill/plugins/huggingface.py

@@ -157,6 +157,22 @@ def _get_unwrapped_teacher(self):
         """Unwrap teacher model (removes FSDP/DDP/DeepSpeed wrapper)."""
         return self.accelerator.unwrap_model(self._teacher_model)
 
+    @contextmanager
+    def _ds_gather(self, params):
+        """Gather DS ZeRO-3 partitioned params; no-op if DeepSpeed disabled.
+
+        The teacher is loaded under an active ``zero.Init`` but not wrapped in a
+        DeepSpeedEngine, so its params have no per-module gather hooks and need an
+        explicit gather around any forward use.
+        """
+        if self.is_deepspeed_enabled:
+            import deepspeed
+
+            with deepspeed.zero.GatheredParameters(list(params), modifier_rank=None):
+                yield
+        else:
+            yield
+
     def compute_loss(self, model, inputs, **kwargs):
         """Store teacher inputs before delegating to parent (which handles liger ctx)."""
         self._ensure_teacher_prepared()
@@ -169,7 +185,7 @@ def compute_kd_loss_func(self, outputs, labels, **kwargs):
         Teacher forward runs here so it is inside the liger identity-lm_head
         context when liger is enabled (ModelOptHFTrainer wraps compute_loss).
         """
-        with torch.no_grad():
+        with torch.no_grad(), self._ds_gather(self._teacher_model.parameters()):
             self._teacher_model.eval()
             teacher_outputs = self._teacher_model(**self._teacher_inputs)
         self._teacher_inputs = None
@@ -193,10 +209,6 @@ def _standard_kd_loss(self, outputs, labels, **kwargs):
         self._last_teacher_outputs = None
         return loss
 
-    def _get_lm_head(self, model):
-        """Resolve lm_head from a model."""
-        return model.lm_head
-
     @contextmanager
     def _liger_identity_lm_head(self):
         """Patch both student+teacher lm_heads to identity."""
@@ -215,22 +227,24 @@ def _liger_identity_lm_head(self):
             teacher_lm_head.forward = teacher_orig
 
     def _sharded_liger_compute(self, fn):
-        """Route fn through sharded DP, gathering both student+teacher lm_head params."""
+        """Delegate student lm_head gather to parent; add teacher lm_head gather on top."""
+        return super()._sharded_liger_compute(self._apply_teacher_gather(fn))
+
+    def _apply_teacher_gather(self, fn):
+        """Wrap fn so the teacher's lm_head params are gathered when it runs."""
         if self.is_fsdp_enabled:
-            return _forward_redirect(
-                self.model,
-                lambda: _forward_redirect(self._teacher_model, fn),
-            )
+            teacher = self._teacher_model
+            return lambda: _forward_redirect(teacher, fn)
         if self.is_deepspeed_enabled:
-            model = self.accelerator.unwrap_model(self.model)
-            teacher = self._get_unwrapped_teacher()
-            student_lm_head = self._get_lm_head(model)
-            teacher_lm_head = self._get_lm_head(teacher)
-            return _forward_redirect(
-                student_lm_head,
-                lambda: _forward_redirect(teacher_lm_head, fn),
-            )
-        return fn()
+            # Teacher is not in the DS engine; gather its lm_head explicitly.
+            teacher_lm_head = self._get_lm_head(self._get_unwrapped_teacher())
+
+            def _wrapped():
+                with self._ds_gather([teacher_lm_head.weight]):
+                    return fn()
+
+            return _wrapped
+        return fn
 
     def _liger_kd_loss(self, outputs, labels, **kwargs):
         """Fused lm_head + JSD for KD."""
@@ -239,13 +253,13 @@ def _liger_kd_loss(self, outputs, labels, **kwargs):
         model = self.accelerator.unwrap_model(self.model)
         teacher = self._get_unwrapped_teacher()
 
-        student_hs = outputs.logits
-        teacher_hs = self._last_teacher_outputs.logits
-        self._last_teacher_outputs = None
-
         student_lm_head = self._get_lm_head(model)
         teacher_lm_head = self._get_lm_head(teacher)
 
+        student_hs = outputs.logits.to(student_lm_head.weight.dtype)  # RMSNorm may upcast to fp32
+        teacher_hs = self._last_teacher_outputs.logits.to(teacher_lm_head.weight.dtype)
+        self._last_teacher_outputs = None
+
         # Causal LM shift
         student_hs = student_hs[..., :-1, :].contiguous().view(-1, student_hs.size(-1))
         teacher_hs = teacher_hs[..., :-1, :].contiguous().view(-1, teacher_hs.size(-1))

modelopt/torch/opt/plugins/transformers.py

@@ -493,8 +493,13 @@ def _prepare_model(self, model):
         """Prepare a model via accelerator (materializes meta-device params, applies sharding).
 
         Uses a dummy optimizer because ``accelerator.prepare`` requires one for FSDP2.
-        Works generically for FSDP2, DDP, and DeepSpeed backends.
+        Works generically for FSDP2, DDP, and DeepSpeed backends. For fully-frozen models
+        under DS ZeRO-3, falls back to inference-mode prep since ZeRO-3 asserts on empty
+        trainable_param_groups; in that case the caller is responsible for gathering
+        ``zero.Init``-partitioned params around forward passes.
         """
+        if self.is_deepspeed_enabled and not any(p.requires_grad for p in model.parameters()):
+            return self.accelerator.prepare_model(model, evaluation_mode=True)
         dummy_optimizer = torch.optim.SGD([next(model.parameters())], lr=0.0)
         model, _ = self.accelerator.prepare(model, dummy_optimizer)
         return model
@@ -712,8 +717,8 @@ def _liger_loss_func(self, outputs, labels, num_items_in_batch=None, **kwargs):
         from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
         model = self.accelerator.unwrap_model(self.model)
-        hidden_states = outputs.logits
         lm_head = self._get_lm_head(model)
+        hidden_states = outputs.logits.to(lm_head.weight.dtype)  # RMSNorm may upcast to fp32
 
         def _compute():
             return LigerForCausalLMLoss(