feat(moe): enable CPU optimizer offload for Kimi-K3 (cuda:nccl,cpu:gloo)

nemo_automodel/components/checkpoint/checkpointing.py

@@ -890,10 +890,20 @@ def load_model(
 
         # MoE adapters return views into model storage; DCP writes safetensors
         # data straight through them and from_hf skips the rebuild.
+        #
+        # ``dequantize_base_checkpoint`` describes the *base* HF checkpoint being
+        # MXFP4-packed, so it only applies when initializing from that base
+        # (is_init_step=True), where the adapter must build packed
+        # ``.weight_packed``/``.weight_scale`` load destinations and dequantize.
+        # Training checkpoints are always saved dequantized (save_model uses
+        # quantization=False), so a resume/restore (is_init_step=False) must read
+        # plain ``.weight`` keys; requesting packed destinations there looks for
+        # ``.weight_packed`` keys the checkpoint does not contain and fails the
+        # DCP load plan.
         state_dict = _maybe_adapt_state_dict_to_hf(
             model_state.model[0],
             state_dict,
-            quantization=self.config.dequantize_base_checkpoint,
+            quantization=bool(self.config.dequantize_base_checkpoint) and is_init_step,
             device_mesh=self.moe_mesh,
         )
 

nemo_automodel/components/distributed/mesh_utils.py

@@ -170,27 +170,70 @@ def _init_named_mesh(
     return device_mesh
 
 
+def _default_pg_has_cpu_backend() -> bool:
+    """Whether the default process group carries a CPU (gloo) co-backend.
+
+    A run that enables CPU parameter/optimizer offload (``CPUOffloadPolicy``)
+    initializes the default process group with a ``cuda:nccl,cpu:gloo`` backend
+    so checkpoint save/load can run DTensor collectives on CPU-resident shards.
+    Detecting that CPU backend lets the per-axis mesh subgroups mirror it. A
+    plain ``nccl`` run reports only ``cuda:nccl`` and keeps NCCL-only subgroups.
+    """
+    if not (dist.is_available() and dist.is_initialized()):
+        return False
+    try:
+        return "cpu:" in dist.get_backend_config()
+    except (RuntimeError, ValueError):
+        return False
+
+
 def _nccl_backend_override(
     axes: tuple[str, ...],
     *,
     device_type: str,
     timeout_minutes: int | None,
-):
-    """Create per-axis NCCL options for DeviceMesh subgroups.
+) -> dict[str, tuple[str, "dist.ProcessGroupNCCL.Options"]] | None:
+    """Create per-axis backend options for DeviceMesh subgroups.
 
     ``init_process_group(timeout=...)`` configures the default process group, but
     ``init_device_mesh`` creates additional per-axis process groups. Without a
     backend override those groups keep PyTorch's default NCCL timeout.
+
+    When the run enabled CPU offload -- signalled by a ``cpu:gloo`` co-backend on
+    the default process group (see :func:`_default_pg_has_cpu_backend`) -- every
+    per-axis subgroup must carry the same ``cuda:nccl,cpu:gloo`` co-backend.
+    Checkpoint save/load then runs DTensor collectives on CPU-resident shards,
+    which an NCCL-only subgroup rejects with ``No backend type associated with
+    device type cpu``. The gloo backend ignores ``ProcessGroupNCCL.Options``, so
+    the NCCL timeout is preserved for the CUDA collectives that use it.
+
+    Args:
+        axes: Mesh axis names to build backend overrides for.
+        device_type: Device type of the mesh; overrides apply only to ``"cuda"``.
+        timeout_minutes: NCCL timeout for the subgroups, or ``None`` to keep the
+            PyTorch default.
+
+    Returns:
+        A mapping from axis name to ``(backend, options)`` for
+        ``init_device_mesh``/``_flatten``/``_unflatten``, or ``None`` when no
+        override is needed (non-CUDA mesh, or CUDA mesh with neither a custom
+        timeout nor a CPU co-backend).
     """
-    if timeout_minutes is None or device_type != "cuda":
+    if device_type != "cuda":
         return None
 
-    timeout = datetime.timedelta(minutes=timeout_minutes)
+    cpu_co_backend = _default_pg_has_cpu_backend()
+    if timeout_minutes is None and not cpu_co_backend:
+        return None
+
+    backend = "cuda:nccl,cpu:gloo" if cpu_co_backend else "nccl"
+    timeout = datetime.timedelta(minutes=timeout_minutes) if timeout_minutes is not None else None
     override = {}
     for axis in axes:
         options = dist.ProcessGroupNCCL.Options()
-        options._timeout = timeout
-        override[axis] = ("nccl", options)
+        if timeout is not None:
+            options._timeout = timeout
+        override[axis] = (backend, options)
     return override
 
 
@@ -425,7 +468,12 @@ def _unflatten_compat(
 ) -> DeviceMesh:
     """Unflatten a mesh with its NCCL timeout, including the PyTorch 2.9 fallback."""
     if hasattr(flat_mesh, "_unflatten"):
-        if timeout_minutes is not None and flat_mesh.device_type == "cuda":
+        # Apply a backend override when the mesh needs a custom NCCL timeout, or
+        # when CPU offload requires the gloo co-backend on the expert subgroups.
+        needs_override = flat_mesh.device_type == "cuda" and (
+            timeout_minutes is not None or _default_pg_has_cpu_backend()
+        )
+        if needs_override:
             return flat_mesh._unflatten(
                 axis,
                 sizes,

nemo_automodel/components/models/kimi_k3/model.py

@@ -1866,8 +1866,17 @@ def meta(*shape: int, tensor_dtype: torch.dtype = dtype) -> torch.Tensor:
 
         emits_hidden_states = getattr(self, "_pp_return_hidden_states", False) is True
         if self.lm_head is not None and not emits_hidden_states:
-            head_dtype = getattr(getattr(self.lm_head, "weight", None), "dtype", dtype)
-            outputs_meta = (meta(microbatch_size, seq_len, text_config.vocab_size, tensor_dtype=head_dtype),)
+            # Logits cross the PP stage boundary in the pipeline compute dtype
+            # (``dtype``, derived from the FSDP mixed-precision activation dtype),
+            # not the lm_head weight's storage dtype. Under fp32-master weights
+            # (``torch_dtype: float32``) with bf16 mixed-precision compute, the head
+            # weight is stored in fp32 but FSDP2 casts it to bf16 for the forward,
+            # so the emitted logits are bf16. lm_head is not in
+            # ``_keep_in_fp32_modules`` and ``compute_lm_head_logits`` runs with
+            # ``fp32_lm_head=False``, so the output always follows the compute dtype.
+            # Keying the meta off the weight dtype trips PipeliningShapeError
+            # ("expected float32 actual bfloat16") at the last stage.
+            outputs_meta = (meta(microbatch_size, seq_len, text_config.vocab_size),)
         elif self.model.norm is not None or block_size is None:
             outputs_meta = (meta(microbatch_size, seq_len, hidden_size),)
         else:

nemo_automodel/components/moe/experts.py

@@ -304,30 +304,36 @@ def forward(
             f"Number of experts must be divisible by ep_size (ep_size={ep_size})"
         )
 
-        # Cast expert weights to the activation dtype so that fp32-stored
-        # parameters (e.g. under fp32 master weights) still work with kernels
-        # (grouped_gemm / torch._grouped_mm) that require matching dtypes with
-        # the (typically bf16) activations. When the weights are already in the
-        # activation dtype these casts are no-ops.
+        # Move expert weights to the activation device and dtype for the compute.
+        #   - dtype: fp32-stored parameters (e.g. under fp32 master weights) must
+        #     match the (typically bf16) activations for grouped_gemm /
+        #     torch._grouped_mm, which require matching operand dtypes.
+        #   - device: under CPUOffloadPolicy the local expert shard is offloaded
+        #     to CPU, and EP-sharded experts are never FSDP-all-gathered back to
+        #     GPU, so the weight must be streamed to the activation device here
+        #     (the expert grouped-matmul is a CUDA kernel). This transient GPU
+        #     copy is exactly the just-in-time streaming CPU offload intends.
+        # When the weights already match ``x`` these calls are no-ops.
         compute_dtype = x.dtype
+        compute_device = x.device
         gate_and_up_projs = (
             self.gate_and_up_projs.to_local() if isinstance(self.gate_and_up_projs, DTensor) else self.gate_and_up_projs
-        ).to(compute_dtype)
+        ).to(device=compute_device, dtype=compute_dtype)
         down_projs = (self.down_projs.to_local() if isinstance(self.down_projs, DTensor) else self.down_projs).to(
-            compute_dtype
+            device=compute_device, dtype=compute_dtype
         )
         gate_up_proj_bias = (
             (
                 self.gate_up_proj_bias.to_local()
                 if isinstance(self.gate_up_proj_bias, DTensor)
                 else self.gate_up_proj_bias
-            ).to(compute_dtype)
+            ).to(device=compute_device, dtype=compute_dtype)
             if self.expert_bias
             else None
         )
         down_proj_bias = (
             (self.down_proj_bias.to_local() if isinstance(self.down_proj_bias, DTensor) else self.down_proj_bias).to(
-                compute_dtype
+                device=compute_device, dtype=compute_dtype
             )
             if self.expert_bias
             else None

nemo_automodel/components/training/signal_handler.py

@@ -36,18 +36,23 @@ def get_device(local_rank: Optional[int] = None) -> torch.device:
         The torch.device ('cuda' for NCCL, 'cpu' for Gloo).
 
     Raises:
-        RuntimeError: If the distributed backend is neither 'nccl' nor 'gloo'.
+        RuntimeError: If the distributed backend exposes neither 'nccl' nor 'gloo'.
     """
-    backend = torch.distributed.get_backend()
-    if backend == "nccl":
+    # ``get_backend()`` returns the plain backend name ("nccl"/"gloo") for a
+    # single-backend group, but a device-typed co-backend group (e.g. CPU offload
+    # runs initialized with "cuda:nccl,cpu:gloo") reports the combined config
+    # string. Match on substring and prefer the CUDA (NCCL) backend when present,
+    # since the caller's collective runs on the default group.
+    backend = str(torch.distributed.get_backend()).lower()
+    if "nccl" in backend:
         if local_rank is None:
             device = torch.device("cuda")
         else:
             device = torch.device(f"cuda:{local_rank}")
-    elif backend == "gloo":
+    elif "gloo" in backend:
         device = torch.device("cpu")
     else:
-        raise RuntimeError
+        raise RuntimeError(f"Unsupported distributed backend {backend!r}; expected 'nccl' or 'gloo'.")
     return device
 
 

tests/unit_tests/components/training/test_signal_handler.py

@@ -51,6 +51,18 @@ def test_get_device_gloo(monkeypatch):
     assert dev.type == "cpu"
 
 
+def test_get_device_cuda_gloo_co_backend(monkeypatch):
+    """
+    A device-typed co-backend group (CPU offload runs use "cuda:nccl,cpu:gloo")
+    must resolve to CUDA, since the NCCL backend runs the collective on cuda.
+    """
+    monkeypatch.setattr(torch.distributed, "get_backend", lambda: "cuda:nccl,cpu:gloo")
+
+    dev = sutils.get_device(local_rank=2)
+    assert dev.type == "cuda"
+    assert dev.index == 2
+
+
 def test_get_device_unknown_backend(monkeypatch):
     """
     An unsupported backend must raise RuntimeError.

tests/unit_tests/distributed/test_mesh_utils.py

@@ -27,6 +27,7 @@
     _create_moe_mesh,
     _init_named_mesh,
     _MeshSpec,
+    _nccl_backend_override,
     _register_flattened_axes,
     _unflatten_compat,
     get_flat_mesh,
@@ -130,6 +131,88 @@ def test_flattened_axes_omit_nccl_timeout_when_unconfigured():
     assert source_mesh._flatten.call_args.kwargs["backend_override"] is None
 
 
+def test_backend_override_uses_gloo_co_backend_when_offload_enabled(monkeypatch):
+    # A CPU (gloo) co-backend on the default PG signals CPU offload; the mesh
+    # subgroups must mirror it so checkpoint save/load can run CPU collectives.
+    monkeypatch.setattr(mesh_utils, "_default_pg_has_cpu_backend", lambda: True)
+
+    override = _nccl_backend_override(
+        (MeshAxisName.PP, MeshAxisName.EP),
+        device_type="cuda",
+        timeout_minutes=30,
+    )
+
+    for axis in (MeshAxisName.PP, MeshAxisName.EP):
+        backend, options = override[axis]
+        assert backend == "cuda:nccl,cpu:gloo"
+        # The NCCL timeout is preserved for the CUDA collectives.
+        assert options._timeout == datetime.timedelta(minutes=30)
+
+
+def test_backend_override_stays_nccl_only_without_offload(monkeypatch):
+    monkeypatch.setattr(mesh_utils, "_default_pg_has_cpu_backend", lambda: False)
+
+    override = _nccl_backend_override(
+        (MeshAxisName.PP,),
+        device_type="cuda",
+        timeout_minutes=30,
+    )
+
+    backend, options = override[MeshAxisName.PP]
+    assert backend == "nccl"
+    assert options._timeout == datetime.timedelta(minutes=30)
+
+
+def test_backend_override_applies_co_backend_without_timeout(monkeypatch):
+    # Offload alone (no custom timeout) must still bind the gloo co-backend.
+    monkeypatch.setattr(mesh_utils, "_default_pg_has_cpu_backend", lambda: True)
+
+    override = _nccl_backend_override(
+        (MeshAxisName.EP,),
+        device_type="cuda",
+        timeout_minutes=None,
+    )
+
+    backend, _ = override[MeshAxisName.EP]
+    assert backend == "cuda:nccl,cpu:gloo"
+
+
+def test_backend_override_none_when_no_timeout_and_no_offload(monkeypatch):
+    monkeypatch.setattr(mesh_utils, "_default_pg_has_cpu_backend", lambda: False)
+
+    assert _nccl_backend_override((MeshAxisName.PP,), device_type="cuda", timeout_minutes=None) is None
+
+
+def test_backend_override_skipped_for_cpu_mesh(monkeypatch):
+    # A non-CUDA mesh never gets an override, even under offload.
+    monkeypatch.setattr(mesh_utils, "_default_pg_has_cpu_backend", lambda: True)
+
+    assert _nccl_backend_override((MeshAxisName.PP,), device_type="cpu", timeout_minutes=30) is None
+
+
+def test_default_pg_has_cpu_backend_false_when_uninitialized(monkeypatch):
+    monkeypatch.setattr(mesh_utils.dist, "is_available", lambda: True)
+    monkeypatch.setattr(mesh_utils.dist, "is_initialized", lambda: False)
+
+    assert mesh_utils._default_pg_has_cpu_backend() is False
+
+
+def test_default_pg_has_cpu_backend_detects_gloo_co_backend(monkeypatch):
+    monkeypatch.setattr(mesh_utils.dist, "is_available", lambda: True)
+    monkeypatch.setattr(mesh_utils.dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(mesh_utils.dist, "get_backend_config", lambda: "cuda:nccl,cpu:gloo")
+
+    assert mesh_utils._default_pg_has_cpu_backend() is True
+
+
+def test_default_pg_has_cpu_backend_false_for_plain_nccl(monkeypatch):
+    monkeypatch.setattr(mesh_utils.dist, "is_available", lambda: True)
+    monkeypatch.setattr(mesh_utils.dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(mesh_utils.dist, "get_backend_config", lambda: "cuda:nccl")
+
+    assert mesh_utils._default_pg_has_cpu_backend() is False
+
+
 def test_fsdp2_forwards_nccl_timeout_to_moe_mesh(monkeypatch):
     device_mesh = Mock()
     moe_mesh = Mock()

tests/unit_tests/models/kimi_k3/test_pipeline_parallel.py

@@ -420,3 +420,23 @@ def test_pipeline_stage_metas_include_block_residual():
 
     assert [tensor.shape for tensor in inputs] == [(2, 4, 32), (8, 1, 32)]
     assert [tensor.shape for tensor in outputs] == [(2, 4, 64)]
+
+
+def test_last_stage_logits_meta_follows_pipeline_dtype_not_weight_dtype():
+    # fp32-master weights (torch_dtype: float32) with bf16 mixed-precision compute:
+    # lm_head.weight is stored in fp32 but the forward emits bf16 logits. The
+    # last-stage output meta must follow the pipeline compute dtype (bf16), not the
+    # fp32 weight storage dtype -- otherwise PP shape inference expects fp32 while
+    # the real output is bf16 and raises PipeliningShapeError at the last stage.
+    model = KimiK3ForCausalLM(_tiny_config(), backend=_torch_backend())
+    assert model.lm_head.weight.dtype == torch.float32
+
+    _, outputs = model.get_pipeline_stage_metas(
+        is_first=True,
+        microbatch_size=2,
+        seq_len=4,
+        dtype=torch.bfloat16,
+    )
+
+    assert outputs[0].shape == (2, 4, 64)
+    assert outputs[0].dtype == torch.bfloat16