Raise when packing is combined with context parallelism

[qgallouedec]

docs/source/distributing_training.md currently recommends packing as a best practice for context parallelism:

Use packing with padding - The default BFD (Best Fit Decreasing) strategy works perfectly:

• Preserves sequence boundaries and maintains training quality

Sequence boundaries are not preserved under context parallelism. Packed documents are kept apart by a block-diagonal attention mask, and context parallelism can only express full causal attention, accelerate's CP hook replaces every layer's mask with is_causal=True. So each packed document attends to every document before it in the same sequence, silently.

Repro

Weight-independent probe: perturb one token in the first packed document and count how many output positions change. Only the first document should be affected. The model uses plain full attention, so packing is the only thing under test.

torchrun --nproc_per_node=2 repro_cp_packing.py

no CP: 50.0% of positions change  (50% means the two documents are correctly separated)
CP=2 : 100.0%  (100% means the second document sees the first)

repro_cp_packing.py

import argparse
import os

import torch
import torch.distributed as dist
from accelerate.big_modeling import _attach_context_parallel_hooks
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.experimental import context_parallel
from transformers import AutoModelForCausalLM, Qwen3Config

DOC = 256  # two packed documents of this length


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dump", help="save the per-position change mask for plotting")
    args = parser.parse_args()

    rank = int(os.environ["RANK"])
    torch.cuda.set_device(rank)
    dist.init_process_group("nccl")
    device = torch.device("cuda", rank)
    mesh = init_device_mesh("cuda", (dist.get_world_size(),), mesh_dim_names=("cp",))

    config = Qwen3Config(vocab_size=1024, hidden_size=256, intermediate_size=512, num_hidden_layers=2,
                         num_attention_heads=4, num_key_value_heads=2, head_dim=64,
                         attn_implementation="sdpa")
    torch.manual_seed(0)
    model = AutoModelForCausalLM.from_config(config).to(device, torch.bfloat16).eval()

    ids_a = torch.randint(0, 1024, (1, 2 * DOC), device=device)
    ids_b = ids_a.clone()
    ids_b[0, 0] = (ids_a[0, 0] + 1) % 1024  # perturb the first token of the first document
    # restarting position_ids is what marks the document boundary
    positions = torch.cat([torch.arange(DOC), torch.arange(DOC)]).unsqueeze(0).to(device)

    def changed(cp):
        def logits(ids):
            kwargs = dict(input_ids=ids, position_ids=positions, use_cache=False)
            if not cp:
                with torch.no_grad():
                    return model(**kwargs).logits.float()
            with torch.no_grad(), context_parallel(mesh, buffers=[ids, positions], buffer_seq_dims=[1, 1]):
                return model(**kwargs).logits.float()

        delta = (logits(ids_a) - logits(ids_b)).abs().amax(-1)[0] > 1e-3
        if not cp:
            return delta
        gathered = [torch.zeros_like(delta) for _ in range(dist.get_world_size())]
        dist.all_gather(gathered, delta)
        return torch.cat(gathered)

    reference = changed(cp=False)
    _attach_context_parallel_hooks(model)
    with_cp = changed(cp=True)

    if rank == 0:
        print(f"no CP: {reference.float().mean():.1%} of positions change  "
              f"(50% means the two documents are correctly separated)")
        print(f"CP={dist.get_world_size()} : {with_cp.float().mean():.1%}  "
              f"(100% means the second document sees the first)")
        if args.dump:
            torch.save({"reference": reference.cpu(), "with_cp": with_cp.cpu(), "doc": DOC}, args.dump)
    dist.destroy_process_group()


if __name__ == "__main__":
    main()

The fix

Raise in SFTTrainer.__init__ when args.packing and parallelism_config.cp_enabled, and correct the doc section. This matches how the other trainers handle context parallelism: GRPOConfig, RLOOConfig and DistillationConfig already raise for cp_size > 1.

Note that padding-free / flash-attention packing is unaffected — this only concerns the context-parallel path, which requires SDPA.

Context

The root cause is in accelerate (its CP hook drops the mask without checking it); that side is addressed in huggingface/accelerate#4177, which fixes the model-level (sliding-window) case. Packing is data-level, so it can only be caught here.

---

Note

Medium Risk
Changes only misconfigured CP+packing setups from silent wrong training to a hard error; valid configs are unchanged, but anyone relying on the old doc guidance must update their SFT config.

Overview
Prevents silent cross-document attention when Ring Attention / context parallelism (cp_size > 1) is used with SFT packing.

SFTTrainer now raises a clear ValueError at init if Accelerate reports cp_enabled and either packing or eval_packing is on. The check uses self.accelerator.parallelism_config (not training args) so it still applies when CP is set only in an Accelerate YAML, and is gated on Accelerate ≥ 1.10.1.

The distributing training docs are updated to match: the Ring Attention example drops packing=True, and best practices now say not to use packing with CP because CP forces full causal attention and cannot preserve packed sequences’ block-diagonal masks.

^{Reviewed by Cursor Bugbot for commit 8cc16d9. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: d748ed2b6e

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit cb37381. Configure here.}

[qgallouedec]

merging without review, feel free to review later if you think I missed something