diff --git a/nemo_automodel/components/checkpoint/checkpointing.py b/nemo_automodel/components/checkpoint/checkpointing.py index 6f659ca11b..1cbfc3d458 100644 --- a/nemo_automodel/components/checkpoint/checkpointing.py +++ b/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, ) diff --git a/nemo_automodel/components/distributed/mesh_utils.py b/nemo_automodel/components/distributed/mesh_utils.py index 3d9ca0e916..af1aab6e0d 100644 --- a/nemo_automodel/components/distributed/mesh_utils.py +++ b/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, diff --git a/nemo_automodel/components/models/kimi_k3/model.py b/nemo_automodel/components/models/kimi_k3/model.py index 9a5050f615..9e04958201 100644 --- a/nemo_automodel/components/models/kimi_k3/model.py +++ b/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: diff --git a/nemo_automodel/components/moe/experts.py b/nemo_automodel/components/moe/experts.py index 74478847e1..dad1063923 100644 --- a/nemo_automodel/components/moe/experts.py +++ b/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 diff --git a/nemo_automodel/components/training/signal_handler.py b/nemo_automodel/components/training/signal_handler.py index 05d4d1ee52..9fb20b31bd 100644 --- a/nemo_automodel/components/training/signal_handler.py +++ b/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 diff --git a/tests/unit_tests/components/training/test_signal_handler.py b/tests/unit_tests/components/training/test_signal_handler.py index f7f71ecb85..607568f574 100644 --- a/tests/unit_tests/components/training/test_signal_handler.py +++ b/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. diff --git a/tests/unit_tests/distributed/test_mesh_utils.py b/tests/unit_tests/distributed/test_mesh_utils.py index be4168c125..3d169d8bba 100644 --- a/tests/unit_tests/distributed/test_mesh_utils.py +++ b/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() diff --git a/tests/unit_tests/models/kimi_k3/test_pipeline_parallel.py b/tests/unit_tests/models/kimi_k3/test_pipeline_parallel.py index 57b0077216..da7af7399f 100644 --- a/tests/unit_tests/models/kimi_k3/test_pipeline_parallel.py +++ b/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 diff --git a/tests/unit_tests/moe/test_experts.py b/tests/unit_tests/moe/test_experts.py index 168e34a658..5d59ab78db 100644 --- a/tests/unit_tests/moe/test_experts.py +++ b/tests/unit_tests/moe/test_experts.py @@ -69,6 +69,54 @@ def moe_config(): ) +class TestGroupedExpertsOffloadDevice: + """GroupedExperts must stream CPU-offloaded expert weights to the activation device. + + Under ``CPUOffloadPolicy`` the sharded (EP-local) expert weights live on CPU and are + never FSDP-all-gathered back to GPU, while activations sit on the compute device. The + forward must move the local weights to ``x.device`` (not merely cast dtype) before the + grouped matmul, otherwise ``torch._grouped_mm`` raises + "mat2 is on cpu, different from other tensors on cuda:N". + """ + + def test_offloaded_weights_moved_to_activation_device(self, moe_config, device): + # fp32-master weights kept on CPU (as CPUOffloadPolicy would); bf16 activations + # on the compute device -> reproduces the offload device split on a CUDA runner. + moe_config.dtype = torch.float32 + backend = BackendConfig(experts="torch_mm") + experts = GroupedExperts(moe_config, backend=backend) # stays on CPU + assert experts.use_torch_mm + with torch.no_grad(): + experts.gate_and_up_projs.normal_(0, 0.02) + experts.down_projs.normal_(0, 0.02) + + num_tokens = 8 + x = torch.randn(num_tokens, moe_config.dim, dtype=torch.bfloat16, device=device) + token_mask = torch.ones(num_tokens, dtype=torch.bool, device=device) + weights = torch.rand(num_tokens, moe_config.n_activated_experts, dtype=torch.bfloat16, device=device) + indices = torch.randint( + 0, moe_config.n_routed_experts, (num_tokens, moe_config.n_activated_experts), device=device + ) + + # Capture the weights handed to the grouped-mm path (avoids the CUDA-only + # torch._grouped_mm so the assertion runs on CPU CI too). + captured = {} + + def _capture(x_arg, token_mask_arg, weights_arg, indices_arg, gate_and_up_projs, down_projs, *rest): + captured["gate_and_up_projs"] = gate_and_up_projs + captured["down_projs"] = down_projs + return torch.zeros(x_arg.shape, dtype=torch.float32, device=x_arg.device) + + with patch.object(experts, "_forward_grouped_mm", _capture): + output = experts(x, token_mask, weights, indices) + + assert captured["gate_and_up_projs"].device == x.device + assert captured["down_projs"].device == x.device + assert captured["gate_and_up_projs"].dtype == x.dtype + assert captured["down_projs"].dtype == x.dtype + assert output.device == x.device + + class TestActivationFunctions: """Test activation functions used in MoE layers."""