transformers 5.16 moves TP to DTensor: all_gather/split are gone, CI red on every PR

[Hotragn]

Summary

transformers 5.16 moved tensor parallelism onto DTensor. transformers.integrations.tensor_parallel is now a back-compat shim over transformers.distributed.tensor_parallel, and the two functions nnsight/modeling/tp/fragments.py depends on — all_gather and split — no longer exist anywhere in transformers.

pyproject.toml pins "transformers" unbounded, so pip install ".[dev]" now resolves 5.16.1 and every PR against 0.8 fails in CI.

CI is currently red for every PR

Same 6 failures on unrelated branches (mine, #699 and #701 — neither touches TP):

FAILED tests/test_tensor_parallel_rules.py::TestStyleCoverage::test_no_upstream_style_is_unaccounted_for
  AssertionError: transformers has parallel styles this version has no rule for:
  ['colwise_rep', 'rowwise_rep']. Add each to SIDES (with the sides that carry a
  shard) or to UNSUPPORTED.

FAILED tests/test_tensor_parallel_rules.py::TestStyleCoverage::test_no_rule_names_a_style_that_no_longer_exists
  AssertionError: rules name parallel styles transformers no longer has:
  ['embedding_colwise'] — probably renamed upstream.

FAILED tests/test_tensor_parallel_rules.py::TestReshardGuard::test_a_value_the_gather_skipped_is_not_split[tensor0-integer position_ids]
FAILED tests/test_tensor_parallel_rules.py::TestReshardGuard::test_a_value_the_gather_skipped_is_not_split[tensor1-a boolean mask]
FAILED tests/test_tensor_parallel_rules.py::TestReshardGuard::test_a_value_the_gather_skipped_is_not_split[tensor2-a 0-dim scalar]
FAILED tests/test_tensor_parallel_rules.py::TestReshardGuard::test_a_real_shard_still_round_trips
  AttributeError: <module 'transformers.integrations.tensor_parallel'> has no attribute 'split'

6 failed, 882 passed, 8 skipped, 1 xfailed

Bisected to the transformers release, not to any nnsight change:

transformers	tests/test_tensor_parallel_rules.py
5.15.1	87 passed
5.16.1	6 failed, 81 passed

The last green run I can see on 0.8 is feat/vllm-graph-taps on 2026-08-24; 5.16 landed after that.

The failing tests understate the problem

The guard tests catch split only because TestReshardGuard monkeypatches it. The imports in fragments.py are function-local, so nothing surfaces until the code actually runs — and _gather's import is never exercised on a CPU runner:

# src/nnsight/modeling/tp/fragments.py:175
def _gather(value, mesh):
    from transformers.integrations.tensor_parallel import all_gather   # <- gone

# src/nnsight/modeling/tp/fragments.py:207
def _reshard(value, mesh):
    from transformers.integrations.tensor_parallel import split        # <- gone

Both now raise:

>>> from transformers.integrations.tensor_parallel import all_gather
ImportError: cannot import name 'all_gather' from 'transformers.integrations.tensor_parallel'
>>> from transformers.integrations.tensor_parallel import split
ImportError: cannot import name 'split' from 'transformers.integrations.tensor_parallel'

So on transformers ≥5.16 a sharded trace should fail at the first fragment handoff, on hardware, with an ImportError from inside a hook. I can't confirm that end to end — no multi-GPU box here — but the imports are unambiguous.

What changed upstream

transformers/integrations/tensor_parallel.py is now:

Backward-compatibility shim for the tensor parallel API.
The canonical implementation lives in transformers.distributed.tensor_parallel.

and the canonical module exports DTensor primitives instead of collectives:

ALL_PARALLEL_STYLES, ColwiseParallel, RowwiseParallel, SequenceParallel, ...
DTensor, Shard, Replicate, Partial, distribute_tensor,
apply_tensor_parallelism, gather_state_dict_for_save, verify_tp_plan

No all_gather, no split, nothing *split*- or *gather*-named besides gather_state_dict_for_save. shard_and_distribute_module is kept but raises:

shard_and_distribute_module is unavailable with the DTensor tensor-parallel loading path.

Style registry deltas:

• added colwise_rep → a ColwiseParallel instance, rowwise_rep → a RowwiseParallel instance (the _rep variants differ in output placement)
• removed embedding_colwise

colwise_rep is the one plan.py already calls out by name — "Llama4Config does exactly this: its plan is colwise_rep, which is in no list and no registry" — it is in the registry now.

Why I'm not sending a patch

The SIDES table is built on "this rank's slice, made whole by an all-gather" / "this rank's term of a sum, made whole by an all-reduce". Under DTensor those handoffs are placements (Shard/Partial/Replicate) rather than explicit collectives, so this reads like a port of modeling/tp/fragments.py rather than a rename — and fragments.py is explicit that the table's inferred entries were the ones that misled before (moe_tp_experts), and that entries were measured on Llama-3.2-3B at tp=4.

I have no multi-GPU machine, so I can't measure which side carries the shard for colwise_rep / rowwise_rep, and guessing into SIDES is exactly what the comments there warn against. Adding them to UNSUPPORTED would turn CI green while preserving today's behaviour (both are currently refused by plan.py, since a style in neither table means "cannot be split") — but that only papers over the missing all_gather/split, so it seemed better to report than to make CI green on a broken path.

Happy to do any of these if you tell me which you want:

1. Unblock CI only — colwise_rep/rowwise_rep → UNSUPPORTED, drop embedding_colwise, and skip TestReshardGuard when split is absent. Green CI, TP still broken on ≥5.16, no false claim of support.
2. Cap the dependency — transformers<5.16 in pyproject.toml plus a MAXIMUM_TRANSFORMERS guard next to the existing MINIMUM_TRANSFORMERS = "5.15.0", so users get UnsupportedTransformersVersion instead of an ImportError from inside a hook.
3. Nothing — if the DTensor port is already in progress on a branch (the TP work in Pipeline parallelism for the vLLM runtime #691/Gather tensor-parallel sharded parameters on read in a trace #677 is yours and @khaiwang's), this is just a heads-up that main-line CI is red meanwhile.

Environment

0.8 @ 96c85fc, transformers 5.16.1 (broken) vs 5.15.1 (fine), torch 2.9.1, Python 3.14.3 locally / 3.12 in CI, CPU.

[Hotragn]

Opened #706, which unblocks CI without touching the part that needs a GPU.

Two things it does:

• Routes both all_gather/split imports through one _collectives() helper that raises UnsupportedTransformersVersion naming what changed and what to pin, instead of a bare ImportError from inside a hook on the first sharded module.
• Skips the six failing tests keyed on the helpers being importable, not on a version number — so a transformers that still has them and adds a style fails them as loudly as before.

Verified on both sides: 5.15.1 → 860 passed / 0 failed; 5.16.1 → 854 passed / 0 failed.

I deliberately left out the two things that are yours to decide:

• No upper version bound. test_the_floor_and_above_are_allowed admits 6.0.0 on purpose; flipping future versions to refuse-by-default reverses that, which shouldn't ride along on a CI fix. (I wrote one first and threw it away.)
• No new SIDES rows for colwise_rep / rowwise_rep. The table's own note says to read the style's _prepare_input_fn/_prepare_output_fn first because the name misled on moe_tp_experts, and the starred rows were measured on Llama-3.2-3B at tp=4. I have no multi-GPU box, and unmeasured rows in that table would defeat the point of it distinguishing measured from inferred.

So this issue stays open for the real work: porting _gather/_reshard to DTensor and re-measuring SIDES against it. Happy to do the mechanical half of that port if you can confirm the handoff points — or to add the upper bound instead if you'd rather declare 5.16 unsupported outright and revisit later.

[JadenFiotto-Kaufman]

Fixed on 0.8 by #709.

Your diagnosis was right on the part you could see, and bisecting it to the
transformers release rather than an nnsight change is what made this quick to
pick up. Thank you.

It turned out to be wider than the two removed helpers. transformers 5.16 changed
three things, and all_gather/split was the least of them.

The plan moved off the modules. Up to 5.15 every sharded module carried
_hf_tp_plan. 5.16 keeps the plan on the model as tp_plan and stamps nothing,
so TPFragments.instrument found nothing sharded: enabled stayed False and
tp_rules stayed empty.

The style's transforms moved into a forward replacement. install_forward
replaces module.forward, which is the same slot nnsight's controller uses, so
installing the controller silently dropped the wrapper.

And all_gather / split were removed, which is the one you found.

Because of the first, _gather/_reshard were never reached on 5.16, so the
ImportError you predicted never fired. What actually happened was
RuntimeError: aten.mm.default got mixed torch.Tensor and DTensor at the first
sharded linear: the module keeping its DTensor weights while losing the code that
matched its inputs to them.

The fixes. Detection now reads model.tp_plan and _device_mesh. The controller
is installed first and transformers re-wraps it. The collectives are written on
DTensor, so they depend on nothing but torch and this class of breakage cannot
recur. Support is 5.16+, but the dependency stays unpinned: only tensor
parallelism needs the floor, and it checks for itself at load rather than holding
every other user to it.

Why nothing caught it: tests/tp was --ignored in CI because it needed GPUs.
transformers TP runs on CPU over gloo, so there is now a 2-rank CPU test that
runs on every PR and would have caught all three changes the day 5.16 landed.

Also fixed while here. Expert parallelism works now, with ep_router and
grouped_gemm leaving UNSUPPORTED after being verified end to end. And four
further bugs two audit passes turned up: a .skip on a sharded module read back
world_size times what was written; backward() through a row-parallel output
deadlocked; an ad-hoc call nested in its own open handoff corrupted it; and values
inside a forward were handed over as raw shards with nothing saying so, which now
warns and has gather/shard helpers plus a docs section.