espnet2.speechlm.model.speechlm.parallel_utils.qwen3.apply_activation_checkpoint_qwen3

source

espnet2.speechlm.model.speechlm.parallel_utils.qwen3.apply_activation_checkpoint_qwen3(model: Module, ac_config: float | List[float] = 0.0, mode: str = 'full', vpp_index: int = 0) → Module

Apply activation checkpointing to transformer layers.

Wraps transformer layers with checkpoint_wrapper for memory savings. Must be applied before torch.compile and FSDP.

When ac_config is a list, it specifies a per-virtual-stage ratio and vpp_index selects which entry to use. When it is a scalar, the same ratio applies to all stages.

nn.Identity placeholder layers (from PP pruning) are skipped.

• Parameters:
  • model – HuggingFace Qwen3 model (possibly PP-pruned with nn.Identity placeholders for non-local layers).
  • ac_config – AC ratio (0.0-1.0) or list of per-virtual-stage ratios. 1.0 = all layers, 0.5 = every other layer, etc.
  • mode – Checkpointing granularity:
    • “full”: wrap the entire transformer layer (default)
    • “moe”: wrap only layer.mlp on MoE layers, skip dense layers.
  • vpp_index – Virtual pipeline stage index on this rank. Used as fallback when model.stage_idx is not set. When ac_config is a list, it is indexed by model.stage_idx (the global virtual stage index).
• Returns: Model with activation checkpointing applied