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

[whn09]

Hi @JingyaHuang — a small follow-up on top of your huggingface#14609, targeting your branch rather than main because the two anchors only exist here. Would you review, and merge it into your branch if you agree with the shape?

The problem

from_pretrained(..., parallel_config=...) shards weights while reading the checkpoint, so enable_parallelism cannot be called afterwards — it raises by design, and the message says exactly why. The loader therefore applies the parallelism itself:

if tp_shard_specs is not None:
    apply_tensor_parallel(model, tp_config, cls._tp_plan, weights_already_sharded=True)

Tensor parallelism, yes. The context-parallel hooks, no. So a ParallelConfig carrying both configs is silently reduced to tensor parallelism alone on this path: no error, no warning, correct output numbers, and every rank redundantly computing the whole sequence. The same holds in _load_dcp_checkpoint.

The change

The context-parallel half of enable_parallelism moves into _apply_context_parallel(config, cp_plan=None) — split out for the same reason you split out _resolve_parallel_config, and the docstring says so — and both sharded-load paths call it before applying tensor parallelism, the order enable_parallelism uses. enable_parallelism itself is now just: resolve, CP, TP.

This is reachable only once ParallelConfig accepts both configs, which is my huggingface#14725. That PR relaxes __post_init__, splits the shared mesh between the two configs, and fixes the Neuron backend's shard index; this one is the loader-side piece it cannot reach.

Why it matters, concretely: MiniMax-H3 has 56 attention heads, so tp_degree caps at 8, and on a 64-core trn2.48xlarge tensor parallelism alone leaves 56 cores idle. At 33B parameters the enable_parallelism route is not an option — every rank would first have to hold the full checkpoint in host memory — so your streaming loader is the only path such a model can take, and without this change it cannot use context parallelism at all. With both PRs, H3 goes from 9.285 s/step on 8 cores to 3.941 s/step on 32 (TP=8 x ulysses=4), a 2.36x speedup.

How to test

I ran this on a trn2.48xlarge. The test is in the PR:

# needs #14725 on top (the guard still rejects TP+CP without it):
git fetch https://github.com/whn09/diffusers tp-cp-compose && git cherry-pick 9c81ac6

python3 -m pytest tests/models/transformers/test_models_transformer_flux.py \
    -k "sharded_load_context_parallel_neuron" -s
# or directly, which is what the test shells out to:
python3 -m torch.distributed.run --nproc_per_node=8 \
    tests/models/transformers/_neuron_sharded_load_worker.py \
    tests.models.transformers.test_models_transformer_flux:make_neuron_sharded_load_spec

Result on my box:

[rank0] tp_degree=2 ulysses_degree=4 context_parallel_hooks=2 output_shape=(1, 16, 4) max_abs_diff=4.1962e-05 max_rel_diff=3.4691e-05
[rank0] PASS: sharded load applied both parallelisms and matches the single-device reference.

Two things worth knowing about how the test is built:

1. Output values cannot detect this bug. A model that skips the CP hooks still returns the right answer — it just does the work redundantly on every rank. So the worker asserts on the structure: that cp_input---* / cp_output---* hooks are registered, and that the attention processors received the ParallelConfig (without which attention runs with no Ulysses all-to-all). It then also compares against a single-device reference read back from the same checkpoint, to catch anything the hooks might break.

2. Negative control. With the one call removed and everything else identical, the worker exits 1 on exactly the intended assertion:

   AssertionError: `from_pretrained(..., parallel_config=...)` applied tensor parallelism but
   registered no context-parallel hooks, so the `context_parallel_config` was silently ignored.
   

Rank 0 writes the checkpoint the other ranks read, so it is single-node, like the other Neuron workers.

Caveats

• The worker is Neuron-only (it needs the "neuron" distributed backend and torchrun, following the _neuron_tp_worker.py convention already in your branch). Porting it to NCCL is essentially the backend string and the device selection — happy to add a CUDA counterpart if you'd rather have one that CI can run.
• I only exercised the safetensors path. The _load_dcp_checkpoint call site is the same one-liner in the same position, but I have not run it; if you have a DCP checkpoint handy that would be a useful second check.
• make_neuron_sharded_load_spec raises num_attention_heads to 8 so the head count survives being divided twice (tp_degree=2 leaves 4 per rank, ulysses_degree=4 splits those into 1 each; ulysses_degree=2 is not available on Neuron, whose all-to-all only accepts group sizes of 4, 8, 16 or multiples of 32). Allow tensor parallelism and context parallelism in one ParallelConfig huggingface/diffusers#14725 adds a make_neuron_hybrid_spec that does the same thing — whichever lands second should reuse the other's.

Separately, the non-persistent-buffer issue I left as a comment on huggingface#14609 is on this same loading path; that one is independent of this PR.