Single-node specdec training broken due to FSDP incompatibility

[benchislett]

Describe the bug

When following the sample example in the EAGLE speculative decoding training on an 8xB200 node, FSDP is used by default but crashes due to some incompatibility / unavailability of FSDP2.

Broken as of #922

Steps/Code to reproduce bug

Follow the default, Llama 3.2-1B EAGLE3 online training from the README:

./launch_train.sh --model meta-llama/Llama-3.2-1B-Instruct\
            --output_dir temp_logs/ \
            --data input_conversations/daring-anteater.jsonl  \
            --num_epochs 1 \
            --eagle_config eagle_config.json --ar_validate_steps -1 --train_bs 1

Crash log:

[rank2]: Traceback (most recent call last):
[rank2]:   File "/root/eagle/ModelOptNew/examples/speculative_decoding/main.py", line 255, in <module>
[rank2]:     train()
[rank2]:   File "/root/eagle/ModelOptNew/examples/speculative_decoding/main.py", line 233, in train
[rank2]:     trainer = EagleTrainerWithAccLog(
[rank2]:   File "/root/eagle/ModelOptNew/.venv/lib/python3.10/site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
[rank2]:     return func(*args, **kwargs)
[rank2]:   File "/root/eagle/ModelOptNew/.venv/lib/python3.10/site-packages/transformers/trainer.py", line 458, in __init__
[rank2]:     self.create_accelerator_and_postprocess()
[rank2]:   File "/root/eagle/ModelOptNew/.venv/lib/python3.10/site-packages/transformers/trainer.py", line 5512, in create_accelerator_and_postprocess
[rank2]:     self.accelerator = Accelerator(**args)
[rank2]:   File "/root/eagle/ModelOptNew/.venv/lib/python3.10/site-packages/accelerate/accelerator.py", line 475, in __init__
[rank2]:     self.parallelism_config._validate_accelerator(self)
[rank2]:   File "/root/eagle/ModelOptNew/.venv/lib/python3.10/site-packages/accelerate/parallelism_config.py", line 364, in _validate_accelerator
[rank2]:     raise ValueError(
[rank2]: ValueError: ParallelismConfig is only compatible DistributedType.FSDP (version 2) or DistributedType.Multi{Device} or DistributedType.DEEPSPEED, but got FSDP.

Expected behavior

Should not crash, should use pure DP (old behaviour) and not FSDP (new behaviour) (is there a difference when all the params can fit on one GPU?)

Who can help?

@yeyu-nvidia @ChenhanYu

System information

• Container used (if applicable): ?
• OS (e.g., Ubuntu 22.04, CentOS 7, Windows 10): Ubuntu 22.04.5 LTS
• CPU architecture (x86_64, aarch64): x86_64
• GPU name (e.g. H100, A100, L40S): NVIDIA B200
• GPU memory size: 179.1 GB
• Number of GPUs: 8
• Library versions (if applicable):
  • Python: 3.10.12
  • ModelOpt version or commit hash: 0.43.0.dev110+g1070d895d
  • CUDA: ?
  • PyTorch: 2.9.1+cu128
  • Transformers: 4.57.6
  • TensorRT-LLM: ?
  • ONNXRuntime: ?
  • TensorRT: ?
• Any other details that may help: ?

[yeyu-nvidia]

Update transformers according to https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/speculative_decoding/requirements.txt should fix this

[ChenhanYu]

Triage Analysis

Field	Value
Classification	Bug
Severity	High
Complexity	Moderate
Auto-fixable	No

Summary

Single-node EAGLE speculative decoding training fails due to FSDP incompatibility. When running the default training script on a single 8-GPU node, the code attempts to use FSDP2 by default, but encounters a version mismatch error where the accelerate library expects DistributedType.FSDP (version 2) but receives the older FSDP type. The issue was introduced in PR #922.

Root Cause / Approach

The issue stems from a mismatch between the transformers and accelerate library versions. The requirements.txt specifies transformers==5.0.0rc1 and accelerate==1.12.0, but the newer accelerate version requires either FSDP2 or other specific distributed types. When FSDP is enabled with fsdp_config.json containing {"fsdp_version":2}, the accelerate library validates that the distributed type is compatible, but the validation fails because the actual distributed type is the older FSDP, not FSDP2.

Suggested Fix

1. Update requirements.txt to pin compatible versions of transformers and accelerate that support FSDP2 properly (likely newer versions). 2. For single-node training where all parameters fit on one GPU, disable FSDP and use standard DistributedDataParallel instead by modifying launch_train.sh to not pass FSDP arguments when running on a single node with sufficient GPU memory. 3. Alternatively, ensure the fsdp_config.json is properly configured for FSDP2 or update the configuration loading logic in main.py to validate and correct incompatibilities before passing to accelerate.

Relevant Files

• examples/speculative_decoding/main.py
• examples/speculative_decoding/launch_train.sh
• examples/speculative_decoding/requirements.txt
• examples/speculative_decoding/fsdp_config.json

---

Auto-triaged by pensieve /magic-triage

[yeyu-nvidia]

I cannot reproduce this error if pip install -r requirements.txt on a single node