Apply the context-parallel hooks on the sharded-load path too

src/diffusers/models/modeling_utils.py

@@ -1655,6 +1655,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
             # was recorded by `_resolve_parallel_config` before loading.
             from ..hooks.tensor_parallel import apply_tensor_parallel
 
+            # A config that also asks for context parallelism gets its hooks here, in the same order as
+            # `enable_parallelism`, which is no longer callable now that the weights are sharded.
+            model._apply_context_parallel(parallel_config)
             apply_tensor_parallel(model, tp_config, cls._tp_plan, weights_already_sharded=True)
         elif parallel_config is not None:
             model.enable_parallelism(config=parallel_config)
@@ -1910,6 +1913,7 @@ def _load_dcp_checkpoint(
         model.register_to_config(_name_or_path=checkpoint_dir)
         model.eval()
 
+        model._apply_context_parallel(parallel_config)
         apply_tensor_parallel(model, tp_config, cls._tp_plan, weights_already_sharded=True)
 
         return model
@@ -2022,11 +2026,6 @@ def enable_parallelism(
             "`enable_parallelism` is an experimental feature. The API may change in the future and breaking changes may be introduced at any time without warning."
         )
 
-        from ..hooks.context_parallel import apply_context_parallel
-        from .attention import AttentionModuleMixin
-        from .attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
-        from .attention_processor import Attention, MochiAttention
-
         if self._parallel_config is not None:
             raise RuntimeError(
                 f"Parallelism is already applied to this {self.__class__.__name__}. `enable_parallelism` cannot be "
@@ -2036,6 +2035,32 @@ def enable_parallelism(
 
         config = self._resolve_parallel_config(config)
 
+        self._apply_context_parallel(config, cp_plan)
+
+        if config.tensor_parallel_config is not None:
+            from ..hooks.tensor_parallel import apply_tensor_parallel
+
+            apply_tensor_parallel(self, config.tensor_parallel_config, self._tp_plan)
+
+    def _apply_context_parallel(
+        self,
+        config: ParallelConfig,
+        cp_plan: dict[str, ContextParallelModelPlan] | None = None,
+    ):
+        """Install the context-parallel hooks `config` asks for, if any, and hand it to the attention processors.
+
+        Split out of `enable_parallelism` for the same reason as `_resolve_parallel_config`: `from_pretrained(...,
+        parallel_config=...)` shards the weights while reading them, so it cannot call `enable_parallelism`
+        afterwards, yet it still has to apply these hooks when the config asks for context parallelism.
+        """
+        from ..hooks.context_parallel import apply_context_parallel
+        from .attention import AttentionModuleMixin
+        from .attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
+        from .attention_processor import Attention, MochiAttention
+
+        if config.context_parallel_config is None:
+            return
+
         attention_classes = (Attention, MochiAttention, AttentionModuleMixin)
 
         if config.context_parallel_config is not None:
@@ -2085,11 +2110,6 @@ def enable_parallelism(
             cp_plan = cp_plan if cp_plan is not None else self._cp_plan
             apply_context_parallel(self, config.context_parallel_config, cp_plan)
 
-        if config.tensor_parallel_config is not None:
-            from ..hooks.tensor_parallel import apply_tensor_parallel
-
-            apply_tensor_parallel(self, config.tensor_parallel_config, self._tp_plan)
-
     @classmethod
     def _load_pretrained_model(
         cls,

tests/models/transformers/_neuron_sharded_load_worker.py

@@ -0,0 +1,170 @@
+# coding=utf-8
+# Copyright 2026 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Generic torchrun worker: assert `from_pretrained(..., parallel_config=...)` installs the context-parallel hooks.
+
+The sharded-load path shards weights while reading the checkpoint, so it cannot call `enable_parallelism` afterwards
+and applies the parallelism itself. Tensor parallelism it did apply; the context-parallel hooks it did not, which is
+silent — the model still returns the right numbers, every rank just redundantly computes the whole sequence. So this
+worker checks both halves: the hooks are registered *and* the output still matches a single-device reference.
+
+Same contract as `_neuron_tp_worker.py`: the model under test is supplied as a `module:function` spec reference on
+the command line, and the referenced factory returns `(model_class, init_dict, inputs)` with CPU tensors.
+
+    torchrun --nproc_per_node=8 _neuron_sharded_load_worker.py \\
+        tests.models.transformers.test_models_transformer_flux:make_neuron_sharded_load_spec
+
+`tp_degree` and `ulysses_degree` are read from `TP_DEGREE` / `ULYSSES_DEGREE` (defaults 2 and 4, whose product is
+the launched world size). `ulysses_degree` cannot be 2 on Neuron: its all-to-all only accepts group sizes of 4, 8,
+16 or multiples of 32.
+
+Rank 0 writes the checkpoint that every rank then reads, so this is single-node only, as the rest of the Neuron
+test workers are.
+
+Exit code 0 means the sharded-load path applied both parallelisms correctly; non-zero means failure.
+"""
+
+import argparse
+import importlib
+import os
+import shutil
+import sys
+import tempfile
+import traceback
+
+
+# Make the in-repo `diffusers` and `tests` packages importable when run via torchrun from an arbitrary CWD.
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..", "src"))
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+
+import torch
+import torch.distributed as dist
+import torch_neuronx  # noqa: F401 — registers torch.neuron
+
+from diffusers import ContextParallelConfig, ParallelConfig, TensorParallelConfig
+
+
+def _context_parallel_hook_names(model):
+    """Every `cp_input---*` / `cp_output---*` hook `apply_context_parallel` registered on `model`."""
+    names = []
+    for module in model.modules():
+        registry = getattr(module, "_diffusers_hook", None)
+        if registry is not None:
+            names.extend(name for name in registry.hooks if name.startswith(("cp_input", "cp_output")))
+    return names
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Neuron sharded-load context-parallel hook worker.")
+    parser.add_argument(
+        "spec",
+        help="`module:function` reference returning (model_class, init_dict, cpu_inputs) for the model under test.",
+    )
+    args = parser.parse_args()
+    module_name, _, fn_name = args.spec.partition(":")
+    model_class, init_dict, inputs = getattr(importlib.import_module(module_name), fn_name)()
+
+    tp_degree = int(os.environ.get("TP_DEGREE", "2"))
+    ulysses_degree = int(os.environ.get("ULYSSES_DEGREE", "4"))
+
+    dist.init_process_group(backend="neuron")
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    device = torch.neuron.current_device()
+
+    if tp_degree * ulysses_degree != world_size:
+        raise ValueError(
+            f"tp_degree ({tp_degree}) x ulysses_degree ({ulysses_degree}) must equal the world size ({world_size})."
+        )
+
+    # Rank 0 writes the checkpoint the sharded loader then reads on every rank. Safetensors, because that is what
+    # the per-rank streaming loader requires.
+    checkpoint_dir = os.environ.get(
+        "SHARDED_LOAD_CHECKPOINT_DIR", os.path.join(tempfile.gettempdir(), "diffusers_neuron_sharded_load_ckpt")
+    )
+    if rank == 0:
+        shutil.rmtree(checkpoint_dir, ignore_errors=True)
+        torch.manual_seed(0)
+        model_class(**init_dict).eval().save_pretrained(checkpoint_dir)
+    dist.barrier()
+
+    # Single-device (unsharded) reference, read back from the same checkpoint so only the loading path differs.
+    ref_output = None
+    if rank == 0:
+        reference = model_class.from_pretrained(checkpoint_dir).eval()
+        with torch.no_grad():
+            ref_output = reference(**inputs, return_dict=False)[0].float().cpu()
+        del reference
+
+    model = model_class.from_pretrained(
+        checkpoint_dir,
+        parallel_config=ParallelConfig(
+            tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree),
+            context_parallel_config=ContextParallelConfig(ulysses_degree=ulysses_degree),
+        ),
+    ).eval()
+    torch.neuron.synchronize()
+
+    # The point of the test: the loader has to install these itself, since `enable_parallelism` is unavailable
+    # once the weights are sharded.
+    hook_names = _context_parallel_hook_names(model)
+    assert hook_names, (
+        "`from_pretrained(..., parallel_config=...)` applied tensor parallelism but registered no context-parallel "
+        "hooks, so the `context_parallel_config` was silently ignored."
+    )
+    processors = [getattr(module, "processor", None) for module in model.modules()]
+    processor_configs = [
+        getattr(p, "_parallel_config", None) for p in processors if p is not None and hasattr(p, "_parallel_config")
+    ]
+    assert processor_configs and all(c is not None for c in processor_configs), (
+        "Context-parallel hooks are registered but the attention processors did not receive the `ParallelConfig`, "
+        "so attention would run without the Ulysses all-to-all."
+    )
+
+    inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
+    with torch.no_grad():
+        output = model(**inputs_on_device, return_dict=False)[0]
+    torch.neuron.synchronize()
+    output = output.float().cpu()
+
+    if rank == 0:
+        assert output.shape == ref_output.shape, f"shape mismatch: {output.shape} vs {ref_output.shape}"
+        assert torch.isfinite(output).all(), "output contains non-finite values"
+        max_abs = (output - ref_output).abs().max().item()
+        denom = ref_output.abs().max().item() + 1e-6
+        print(
+            f"[rank0] tp_degree={tp_degree} ulysses_degree={ulysses_degree} "
+            f"context_parallel_hooks={len(hook_names)} output_shape={tuple(output.shape)} "
+            f"max_abs_diff={max_abs:.4e} max_rel_diff={max_abs / denom:.4e}"
+        )
+        # Neuron runs matmuls in bf16 internally, so compare with a bf16-level tolerance, as `_neuron_tp_worker`
+        # does. A wrong shard plan or a mis-ordered mesh produces grossly different output and is caught well
+        # inside this bound.
+        torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=2e-2)
+        print("[rank0] PASS: sharded load applied both parallelisms and matches the single-device reference.")
+
+    dist.barrier()
+    if rank == 0:
+        shutil.rmtree(checkpoint_dir, ignore_errors=True)
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception:
+        traceback.print_exc()
+        # Ensure a non-zero exit so the launching pytest sees the failure.
+        os._exit(1)

tests/models/transformers/test_models_transformer_flux.py

@@ -281,6 +281,41 @@ def make_neuron_tp_spec():
     return FluxTransformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu")
 
 
+def make_neuron_sharded_load_spec():
+    """Model spec consumed by the sharded-load worker (`_neuron_sharded_load_worker.py`).
+
+    Same contract as `make_neuron_tp_spec`, but `num_attention_heads` is raised to 8 so the head count survives
+    being divided twice: `tp_degree=2` leaves 4 heads per rank and `ulysses_degree=4` splits those into 1 each.
+    (`ulysses_degree` cannot be 2 on Neuron, whose all-to-all only accepts group sizes of 4, 8, 16 or multiples
+    of 32.)
+    """
+    config = FluxTransformerTesterConfig()
+    init_dict = config.get_init_dict() | {"num_attention_heads": 8}
+    return FluxTransformer2DModel, init_dict, config.get_dummy_inputs(device="cpu")
+
+
+@is_tensor_parallel
+@require_torch_neuron
+class TestFluxTransformerShardedLoadNeuron:
+    """`from_pretrained(..., parallel_config=...)` with both parallelisms, on AWS Neuron.
+
+    Launched the same way as `TestFluxTransformerTensorParallelNeuron`, at `tp_degree=2 x ulysses_degree=4`, i.e. 8
+    ranks. The worker asserts that the sharded-load path registered the context-parallel hooks — which numbers alone
+    cannot detect, since a model that silently skips them still returns the right answer — and that the output still
+    matches a single-device reference read back from the same checkpoint.
+    """
+
+    def test_sharded_load_context_parallel_neuron(self):
+        worker = os.path.join(os.path.dirname(__file__), "_neuron_sharded_load_worker.py")
+        spec = "tests.models.transformers.test_models_transformer_flux:make_neuron_sharded_load_spec"
+        cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=8", worker, spec]
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        assert result.returncode == 0, (
+            f"Neuron sharded-load worker failed (exit {result.returncode}).\n"
+            f"--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}"
+        )
+
+
 @is_tensor_parallel
 @require_torch_neuron
 class TestFluxTransformerTensorParallelNeuron: