espnet2.speechlm.model.speechlm.parallel_utils.qwen3.parallelize_qwen3_hf

source

espnet2.speechlm.model.speechlm.parallel_utils.qwen3.parallelize_qwen3_hf(model: Module, parallel_dims: ParallelDims, titan_config: Dict[str, Any], vpp_index: int = 0) → Module

Apply parallelization to HuggingFace Qwen3 model.

Order: AC -> torch.compile -> FSDP (following TorchTitan’s convention)

• Parameters:

  • model – HuggingFace Qwen3 model (possibly wrapped with multimodal components)

  • parallel_dims – TorchTitan ParallelDims object with device meshes

  • titan_config –

    Configuration dict containing:

    • activation_checkpoint: AC ratio 0.0-1.0 (default: 0.0).

    1.0 = all layers, 0.5 = every other layer.

    • compile: Whether to enable torch.compile (default: false)
    • compile_mode: Compile mode (default: “default”)
    • mixed_precision_param: Parameter dtype (default: “bfloat16”)
    • mixed_precision_reduce: Reduce dtype (default: “float32”)
    • reshard_after_forward: Whether to reshard params after forward (default: true). true saves memory, false is faster.
• Returns: Parallelized model