System Info
- transformers 5.17.0.dev0 (main 41f519e)
- accelerate 1.14.0
- torch 2.13.0+cu130
- python 3.13.13
- 2x H100 80GB
Who can help?
@SunMarc
Information
Tasks
Reproduction
Under FSDP2 + context parallelism, Trainer.predict returns predictions with cp_size times too many rows, and the unlabeled path silently skips the context-parallel context entirely.
import torch
from torch.utils.data import Dataset
from transformers import AutoModelForCausalLM, Qwen3Config, Trainer, TrainingArguments
ROWS, SEQ, VOCAB = 8, 512, 1024
class Tokens(Dataset):
def __init__(self, labeled):
self.ids = torch.randint(0, VOCAB, (ROWS, SEQ), generator=torch.Generator().manual_seed(0))
self.labeled = labeled
def __len__(self):
return len(self.ids)
def __getitem__(self, index):
item = {"input_ids": self.ids[index], "attention_mask": torch.ones(SEQ, dtype=torch.long)}
if self.labeled:
item["labels"] = self.ids[index]
return item
def collate(features):
return {key: torch.stack([f[key] for f in features]) for key in features[0]}
config = Qwen3Config(
vocab_size=VOCAB,
hidden_size=128,
intermediate_size=256,
num_hidden_layers=2,
num_attention_heads=4,
num_key_value_heads=2,
max_position_embeddings=SEQ,
attn_implementation="sdpa",
)
model = AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16)
args = TrainingArguments(
output_dir="out",
max_steps=1,
learning_rate=0.0,
per_device_train_batch_size=1,
per_device_eval_batch_size=2,
bf16=True,
save_strategy="no",
report_to=[],
)
trainer = Trainer(model=model, args=args, data_collator=collate, train_dataset=Tokens(labeled=True))
trainer.train() # FSDP2 rejects prepare(model) without an optimizer, so train first
for labeled in (True, False):
predictions = trainer.predict(Tokens(labeled=labeled)).predictions
if trainer.accelerator.is_main_process:
print(f"cp_size={trainer.get_cp_size()} labeled={labeled} predictions={tuple(predictions.shape)}", flush=True)# cp_size=2
accelerate launch --num_processes 2 --use_fsdp --fsdp_version 2 --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--use_parallelism_config --parallelism_config_cp_size 2 repro_cp_predictions.py
# control
accelerate launch --num_processes 1 --use_fsdp --fsdp_version 2 \
--fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP repro_cp_predictions.py
8 rows of 512 tokens:
cp_size=2 labeled=True predictions=(16, 256, 1024)
cp_size=2 labeled=False predictions=(16, 512, 1024)
cp_size=1 labeled=True predictions=(8, 512, 1024)
cp_size=1 labeled=False predictions=(8, 512, 1024)
Two separate things.
Rows are duplicated. The sampler shards over the dp dimension only, so the cp_size ranks of a group all hold the same batch, while gather_function gathers over the whole world, so every row lands once per CP rank. eval_loss is unaffected, duplicating a value doesn't move its mean, which is why #48167 reads clean. Anything going through compute_metrics does not. This one isn't specific to CP: tp_size > 1 shares a batch the same way.
The unlabeled branch isn't sharded. prediction_step only enters the context for the labeled path:
https://github.com/huggingface/transformers/blob/41f519e741/src/transformers/trainer.py#L3143-L3148
The else at L3160 calls the model directly, which is the 512 above: every CP rank recomputes the full sequence. At the lengths CP exists for that's an OOM where training is fine.
The 256 in the labeled row is expected, the context shards the buffers in place and nothing reassembles the logits. Gathering a [batch, seq, vocab] tensor back together would defeat the point of CP, so refusing or warning when predictions are requested under CP seems better than returning them sharded.
Expected behavior
predict under CP returns one row per dataset row, and the unlabeled path shards the sequence like the labeled one, or the Trainer says predictions aren't supported under CP instead of returning wrong shapes.
System Info
Who can help?
@SunMarc
Information
Tasks
Reproduction
Under FSDP2 + context parallelism,
Trainer.predictreturns predictions withcp_sizetimes too many rows, and the unlabeled path silently skips the context-parallel context entirely.8 rows of 512 tokens:
Two separate things.
Rows are duplicated. The sampler shards over the dp dimension only, so the
cp_sizeranks of a group all hold the same batch, whilegather_functiongathers over the whole world, so every row lands once per CP rank.eval_lossis unaffected, duplicating a value doesn't move its mean, which is why #48167 reads clean. Anything going throughcompute_metricsdoes not. This one isn't specific to CP: tp_size > 1 shares a batch the same way.The unlabeled branch isn't sharded.
prediction_steponly enters the context for the labeled path:https://github.com/huggingface/transformers/blob/41f519e741/src/transformers/trainer.py#L3143-L3148
The
elseat L3160 calls the model directly, which is the 512 above: every CP rank recomputes the full sequence. At the lengths CP exists for that's an OOM where training is fine.The 256 in the labeled row is expected, the context shards the buffers in place and nothing reassembles the logits. Gathering a [batch, seq, vocab] tensor back together would defeat the point of CP, so refusing or warning when predictions are requested under CP seems better than returning them sharded.
Expected behavior
predictunder CP returns one row per dataset row, and the unlabeled path shards the sequence like the labeled one, or the Trainer says predictions aren't supported under CP instead of returning wrong shapes.