transformers
cb33194a - Do not wrap the model in FSDP for expert-parallel token dispatch (#49383)

Commit
7 days ago
Do not wrap the model in FSDP for expert-parallel token dispatch (#49383) Do not wrap the model in FSDP for expert-parallel token dispatch when fsdp_size == 1 With fsdp_size == 1 the mesh constraints force ep_size == tp_size, hence efsdp_size == 1: there are no expert replicas to shard or reduce, so the forced fully_shard added nothing but overhead (size-1 FSDP wrapping of every module, synced_gpus in generate, DCP-based state dict gathering). Apply FSDP only when fsdp_size > 1 and route the fsdp_size == 1 save path through the regular DTensor full_tensor() gather.
Author
Parents
Loading