# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from collections import defaultdict import torch import torch.nn as nn from torch.distributed._composable.fsdp import fully_shard from torch.distributed._composable.replicate import replicate from torch.distributed._tensor import Replicate, Shard from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper as ptd_checkpoint_wrapper from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor.parallel import ( ColwiseParallel, PrepareModuleInput, RowwiseParallel, SequenceParallel, parallelize_module, ) from cosmos_predict2._src.imaginaire.utils import log as logger from cosmos_predict2._src.reason1.configs.default.model_config import ActivationCheckpointConfig from cosmos_predict2._src.reason1.configs.default.model_config import FSDP2ModelConfig as JobConfig # from torchtitan.config_manager import TORCH_DTYPE_MAP # from torchtitan.logging import logger from cosmos_predict2._src.reason1.parallelisms.parallel_dims import ParallelDims TORCH_DTYPE_MAP = { "float16": torch.float16, "float32": torch.float32, "bfloat16": torch.bfloat16, } def parallelize_qwen( model: nn.Module, world_mesh: DeviceMesh, parallel_dims: ParallelDims, job_config: JobConfig, ): """ Apply tensor parallelism, activation checkpointing, torch.compile, and data parallelism to the model. NOTE: The passed-in model preferably should be on meta device. Otherwise, the model must fit on GPU or CPU memory. """ if parallel_dims.tp_enabled: if job_config.experimental.enable_async_tensor_parallel and not job_config.training.compile: raise RuntimeError("Async TP requires --training.compile") apply_tp( model, world_mesh["tp"], loss_parallel=parallel_dims.loss_parallel_enabled, enable_float8=job_config.float8.enable_float8_linear, enable_async_tp=job_config.experimental.enable_async_tensor_parallel, ) if job_config.activation_checkpoint.mode != "none": apply_ac(model, job_config.activation_checkpoint) # turn on per-TransformerBlock compile after AC wrapping and before FSDP if job_config.training.compile: apply_compile(model) if ( parallel_dims.dp_shard_enabled or parallel_dims.cp_enabled ): # apply FSDP or HSDP, potentially with Context Parallel if parallel_dims.dp_replicate_enabled: dp_mesh_dim_names = ("dp_replicate", "dp_shard_cp") else: dp_mesh_dim_names = ("dp_shard_cp",) apply_fsdp( model, world_mesh[tuple(dp_mesh_dim_names)], ) if parallel_dims.dp_replicate_enabled: logger.info("Applied HSDP to the model") else: logger.info("Applied FSDP to the model") if parallel_dims.cp_enabled: logger.info("Applied Context Parallel to the model") if job_config.training.enable_cpu_offload: logger.info("Applied CPU Offloading to the model") elif parallel_dims.dp_replicate_enabled: if world_mesh.ndim > 1: raise RuntimeError("DDP has not supported > 1D parallelism") apply_ddp( model, world_mesh, enable_compile=job_config.training.compile, enable_compiled_autograd=job_config.experimental.enable_compiled_autograd, ) def apply_tp( model: nn.Module, tp_mesh: DeviceMesh, loss_parallel: bool, enable_float8: bool, enable_async_tp: bool, ): """Apply tensor parallelism.""" # 1. Parallelize the embedding and shard its outputs (which are the first # transformer block's inputs) # 2. Parallelize the root norm layer over the sequence dim # 3. Parallelize the final linear output layer parallelize_module( model, tp_mesh, { "model.embed_tokens": RowwiseParallel( input_layouts=Replicate(), output_layouts=Shard(1), use_local_output=False, # Output Dtensor ), "model.norm": SequenceParallel(), "lm_head": ColwiseParallel( input_layouts=Shard(1), output_layouts=Shard(-1) if loss_parallel else Replicate(), use_local_output=not loss_parallel, ), }, ) # Parallel styles used for transformer block linear weights and their # inputs may be different for float8 linears if enable_float8: # add a check here to enforce supported float8 all-gather configurations from torchao.float8.float8_tensor_parallel import ( Float8ColwiseParallel, Float8RowwiseParallel, PrepareFloat8ModuleInput, ) rowwise_parallel, colwise_parallel, prepare_module_input = ( Float8RowwiseParallel, Float8ColwiseParallel, PrepareFloat8ModuleInput, ) else: rowwise_parallel, colwise_parallel, prepare_module_input = ( RowwiseParallel, ColwiseParallel, PrepareModuleInput, ) # Apply tensor + sequence parallelism to every transformer block # NOTE: At the cost of model code change, we can accelerate Sequence Parallel # by folding (and unfolding) the batch dimension and the sequence dimension. # Examples can be found at https://github.com/pytorch/torchtitan/pull/437 for transformer_block in model.model.layers: layer_plan = { "attention_norm": SequenceParallel(), "attention": prepare_module_input( input_layouts=( Shard(1), # hidden_states None, # attention_mask None, # position_ids None, # past_key_value None, # output_attentions None, # use_cache None, # cache_position None, # position_embeddings ), desired_input_layouts=( Replicate(), None, # attention_mask None, # position_ids None, # past_key_value None, # output_attentions None, # use_cache None, # cache_position None, # position_embeddings), ), ), "attention.wq": colwise_parallel(), "attention.wk": colwise_parallel(), "attention.wv": colwise_parallel(), "attention.wo": rowwise_parallel(output_layouts=Shard(1)), "ffn_norm": SequenceParallel(), "feed_forward": prepare_module_input( input_layouts=(Shard(1),), desired_input_layouts=(Replicate(),), ), "feed_forward.w1": colwise_parallel(), "feed_forward.w2": rowwise_parallel(output_layouts=Shard(1)), "feed_forward.w3": colwise_parallel(), } # map the name from llama to qwen names_mapping = { "attention_norm": "input_layernorm", "attention": "self_attn", "attention.wq": "self_attn.q_proj", "attention.wk": "self_attn.k_proj", "attention.wv": "self_attn.v_proj", "attention.wo": "self_attn.o_proj", "ffn_norm": "post_attention_layernorm", # Norm after attention, before feed_forward "feed_forward": "mlp", "feed_forward.w1": "mlp.gate_proj", "feed_forward.w2": "mlp.down_proj", "feed_forward.w3": "mlp.up_proj", } new_layer_plan = {} for key, value in layer_plan.items(): new_layer_plan[names_mapping[key]] = value del layer_plan layer_plan = new_layer_plan parallelize_module( module=transformer_block, device_mesh=tp_mesh, parallelize_plan=layer_plan, ) if enable_async_tp: from torch.distributed._symmetric_memory import enable_symm_mem_for_group torch._inductor.config._micro_pipeline_tp = True enable_symm_mem_for_group(tp_mesh.get_group().group_name) logger.info( f"Applied {'Float8 ' if enable_float8 else ''}{'Async ' if enable_async_tp else ''}" "Tensor Parallelism to the model" ) # for selective op activation checkpointing _save_list = { torch.ops.aten.mm.default, torch.ops.aten._scaled_dot_product_efficient_attention.default, torch.ops.aten._scaled_dot_product_flash_attention.default, torch.ops._c10d_functional.reduce_scatter_tensor.default, # for low precision training, it's useful to always save # the result of max, since the absolute maximum is # used to compute the scaling factor for quantization. torch.ops.aten.max.default, } def _apply_ac_to_transformer_block(module: nn.Module, ac_config): valid_ac_modes = ("full", "selective") if ac_config.mode not in valid_ac_modes: raise ValueError(f"Invalid AC mode: {ac_config.mode}. Valid modes: {valid_ac_modes}") if ac_config.mode == "full": return ptd_checkpoint_wrapper(module, preserve_rng_state=False) assert ac_config.mode == "selective", f"{ac_config.mode}" use_op_sac = ac_config.selective_ac_option == "op" use_layer_sac = ac_config.selective_ac_option.isdigit() # print(f"use_op_sac: {use_op_sac}, use_layer_sac: {use_layer_sac}") if not use_op_sac and not use_layer_sac: raise ValueError( f"Invalid selective AC option: {ac_config.selective_ac_option}. " f"Valid options: 'op' or a positive int representing layer frequency" ) if use_op_sac: from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts def _get_custom_policy(meta): def _custom_policy(ctx, func, *args, **kwargs): mode = "recompute" if ctx.is_recompute else "forward" mm_count_key = f"{mode}_mm_count" if func == torch.ops.aten.mm.default: meta[mm_count_key] += 1 # Saves output of all compute ops, except every second mm to_save = func in _save_list and not (func == torch.ops.aten.mm.default and meta[mm_count_key] % 2 == 0) return CheckpointPolicy.MUST_SAVE if to_save else CheckpointPolicy.PREFER_RECOMPUTE return _custom_policy def selective_checkpointing_context_fn(): meta = defaultdict(int) return create_selective_checkpoint_contexts(_get_custom_policy(meta)) return ptd_checkpoint_wrapper( module, # wrapped_forward, context_fn=selective_checkpointing_context_fn, preserve_rng_state=False, ) elif use_layer_sac: # Checkpoint every `ac_freq` of the modules passed to this function ac_freq = int(ac_config.selective_ac_option) ptd_checkpoint_wrapper.__dict__.setdefault("_count", 0) ptd_checkpoint_wrapper._count += 1 if not ac_freq or ptd_checkpoint_wrapper._count % ac_freq == 0: return ptd_checkpoint_wrapper( module, # wrapped_forward, preserve_rng_state=False, ) else: return module def apply_ac(model: nn.Module, ac_config: ActivationCheckpointConfig): """Apply activation checkpointing to the model.""" # model.model is Qwen2_5_VLModel if ( ("vision" == ac_config.models or "vlm" == ac_config.models) and hasattr(model, "visual") and model.visual is not None ): for layer_id, block in model.visual.blocks.named_children(): block = ptd_checkpoint_wrapper(block, preserve_rng_state=False) model.visual.blocks.register_module(layer_id, block) if "llm" == ac_config.models or "vlm" == ac_config.models: for layer_id, transformer_block in model.model.layers.named_children(): transformer_block = _apply_ac_to_transformer_block(transformer_block, ac_config) model.model.layers.register_module(layer_id, transformer_block) logger.info(f"Applied {ac_config.mode} activation checkpointing to the model") def apply_compile(model: nn.Module): """ Apply torch.compile to each TransformerBlock, which makes compilation efficient due to repeated structure. Alternatively one can compile the whole model (after applying DP). """ for layer_id, transformer_block in model.layers.named_children(): transformer_block = torch.compile(transformer_block, fullgraph=True) model.layers.register_module(layer_id, transformer_block) logger.info("Compiling each TransformerBlock with torch.compile") def apply_fsdp( model: nn.Module, dp_mesh: DeviceMesh, ): """ Apply data parallelism (via FSDP2) to the model. Args: model (nn.Module): The model to apply data parallelism to. dp_mesh (DeviceMesh): The device mesh to use for data parallelism. """ if model.visual is not None: for layer_id, block in enumerate(model.visual.blocks): fully_shard(block, mesh=dp_mesh) for layer_id, transformer_block in enumerate(model.model.layers): fully_shard( transformer_block, mesh=dp_mesh, ) fully_shard(model, mesh=dp_mesh) def apply_ddp( model: nn.Module, dp_mesh: DeviceMesh, enable_compile: bool, enable_compiled_autograd: bool, ): if enable_compile: if enable_compiled_autograd: torch._dynamo.config.optimize_ddp = "python_reducer_without_compiled_forward" else: torch._dynamo.config.optimize_ddp = "ddp_optimizer" replicate(model, device_mesh=dp_mesh, bucket_cap_mb=100) logger.info("Applied DDP to the model")