espnet2.speechlm.model.speechlm.parallel_utils.qwen3.apply_fsdp_qwen3

source

espnet2.speechlm.model.speechlm.parallel_utils.qwen3.apply_fsdp_qwen3(model: Module, parallel_dims: ParallelDims, titan_config: Dict[str, Any]) → Module

Apply FSDP2 to HuggingFace Qwen3 model structure.

Moves modules from CPU to GPU one FSDP unit at a time, then shards immediately. This avoids materializing the full model on every GPU (which would waste ~60GB for a 30B model). Peak GPU memory during init is ~1 transformer layer instead of the entire model.

Tolerates pruned PP stage models where some modules (embed_tokens, lm_head, norm, stream_emb, multimodal_io_dict, adaptor) may be None.

• Parameters:
  • model – HuggingFace Qwen3 model (on CPU or GPU) to wrap with FSDP. May be a full model or a pruned PP stage.
  • parallel_dims – TorchTitan ParallelDims with device meshes
  • titan_config – Configuration dict
• Returns: FSDP-wrapped model (on GPU, sharded)