diff --git a/.file_mapping.json b/.file_mapping.json index 503fa929..5b0e65ba 100644 --- a/.file_mapping.json +++ b/.file_mapping.json @@ -1,16 +1,15 @@ { - "_source_commit": "c2c7152e59e4090a61822dc22a1691e7afb44702-dirty", - "_dest_commit": "463ac142b47c637098a7e86d0631252f82b65418", - "_generated_at": "2026-07-07T05:54:17Z", + "_source_commit": "7392c4b59ab7c1f0a94e59fbc4a75a0f3684b66f-dirty", + "_dest_commit": "3d9c0878fd0dde76eac98161aed0493d85a036fd", + "_generated_at": "2026-07-11T05:50:38Z", "files": { "imaginaire/__init__.py": "cosmos_framework/__init__.py", "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py", "imaginaire/attention/backends.py": "cosmos_framework/model/attention/backends.py", - "imaginaire/attention/benchmarks/benchmark_fmha_head_sharding.py": "cosmos_framework/model/attention/benchmarks/benchmark_fmha_head_sharding.py", + "imaginaire/attention/benchmarks/benchmark_fmha.py": "cosmos_framework/model/attention/benchmarks/benchmark_fmha.py", "imaginaire/attention/checks.py": "cosmos_framework/model/attention/checks.py", "imaginaire/attention/cudnn/__init__.py": "cosmos_framework/model/attention/cudnn/__init__.py", "imaginaire/attention/cudnn/checks.py": "cosmos_framework/model/attention/cudnn/checks.py", - "imaginaire/attention/cudnn/cudnn_forward.py": "cosmos_framework/model/attention/cudnn/cudnn_forward.py", "imaginaire/attention/cudnn/functions.py": "cosmos_framework/model/attention/cudnn/functions.py", "imaginaire/attention/cudnn/meta.py": "cosmos_framework/model/attention/cudnn/meta.py", "imaginaire/attention/cudnn/stubs.py": "cosmos_framework/model/attention/cudnn/stubs.py", @@ -88,6 +87,7 @@ "imaginaire/flops/__init__.py": "cosmos_framework/tools/flops/__init__.py", "imaginaire/flops/omni_mot.py": "cosmos_framework/tools/flops/omni_mot.py", "imaginaire/flops/qwen3_vl.py": "cosmos_framework/tools/flops/qwen3_vl.py", + "imaginaire/flops/test_qwen3_vl.py": "cosmos_framework/tools/flops/test_qwen3_vl.py", "imaginaire/flops/wan_vae.py": "cosmos_framework/tools/flops/wan_vae.py", "imaginaire/functional/batch_ops.py": "cosmos_framework/utils/functional/batch_ops.py", "imaginaire/functional/lr_scheduler.py": "cosmos_framework/utils/functional/lr_scheduler.py", @@ -132,6 +132,7 @@ "imaginaire/utils/easy_io/handlers/csv_handler.py": "cosmos_framework/utils/easy_io/handlers/csv_handler.py", "imaginaire/utils/easy_io/handlers/gzip_handler.py": "cosmos_framework/utils/easy_io/handlers/gzip_handler.py", "imaginaire/utils/easy_io/handlers/imageio_video_handler.py": "cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py", + "imaginaire/utils/easy_io/handlers/imageio_video_handler_test.py": "cosmos_framework/utils/easy_io/handlers/imageio_video_handler_test.py", "imaginaire/utils/easy_io/handlers/json_handler.py": "cosmos_framework/utils/easy_io/handlers/json_handler.py", "imaginaire/utils/easy_io/handlers/jsonl_handler.py": "cosmos_framework/utils/easy_io/handlers/jsonl_handler.py", "imaginaire/utils/easy_io/handlers/np_handler.py": "cosmos_framework/utils/easy_io/handlers/np_handler.py", @@ -405,7 +406,9 @@ "projects/cosmos3/cosmos3/tokenizers/wan2pt2_vae_4x16x16.py": "cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py", "projects/cosmos3/cosmos3/upsampler/__init__.py": "cosmos_framework/model/generator/upsampler/__init__.py", "projects/cosmos3/cosmos3/upsampler/prompts.py": "cosmos_framework/model/generator/upsampler/prompts.py", + "projects/cosmos3/cosmos3/utils/aux_optimizer_utils.py": "cosmos_framework/utils/generator/aux_optimizer_utils.py", "projects/cosmos3/cosmos3/utils/data_utils.py": "cosmos_framework/utils/generator/data_utils.py", + "projects/cosmos3/cosmos3/utils/dion2_with_aux_adamw.py": "cosmos_framework/utils/generator/dion2_with_aux_adamw.py", "projects/cosmos3/cosmos3/utils/dtensor_helper.py": "cosmos_framework/utils/generator/dtensor_helper.py", "projects/cosmos3/cosmos3/utils/flash_attn.py": "cosmos_framework/utils/generator/flash_attn.py", "projects/cosmos3/cosmos3/utils/fused_adam.py": "cosmos_framework/utils/generator/fused_adam.py", @@ -414,6 +417,7 @@ "projects/cosmos3/cosmos3/utils/model_loader.py": "cosmos_framework/utils/generator/model_loader.py", "projects/cosmos3/cosmos3/utils/model_weights_stats.py": "cosmos_framework/utils/generator/model_weights_stats.py", "projects/cosmos3/cosmos3/utils/monkey_patch.py": "cosmos_framework/utils/generator/monkey_patch.py", + "projects/cosmos3/cosmos3/utils/muon_with_aux_adamw.py": "cosmos_framework/utils/generator/muon_with_aux_adamw.py", "projects/cosmos3/cosmos3/utils/optimizer.py": "cosmos_framework/utils/generator/optimizer.py", "projects/cosmos3/cosmos3/utils/parallelism.py": "cosmos_framework/utils/generator/parallelism.py", "projects/cosmos3/cosmos3/utils/rand_state.py": "cosmos_framework/utils/generator/rand_state.py", @@ -423,9 +427,11 @@ "projects/cosmos3/cosmos3/utils/reasoner/flop_calculator.py": "cosmos_framework/utils/generator/reasoner/flop_calculator.py", "projects/cosmos3/cosmos3/utils/reasoner/pretrained_models_downloader.py": "cosmos_framework/utils/generator/reasoner/pretrained_models_downloader.py", "projects/cosmos3/cosmos3/utils/video_preprocess.py": "cosmos_framework/utils/generator/video_preprocess.py", + "projects/cosmos3/tokenizer/checkpoint_io.py": "cosmos_framework/model/tokenizer/checkpoint_io.py", "projects/cosmos3/tokenizer/evaluation/metric.py": "cosmos_framework/model/tokenizer/evaluation/metric.py", "projects/cosmos3/tokenizer/evaluation/reconstruction_metrics.py": "cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py", "projects/cosmos3/tokenizer/models/__init__.py": "cosmos_framework/model/tokenizer/models/__init__.py", + "projects/cosmos3/tokenizer/models/architecture.py": "cosmos_framework/model/tokenizer/models/architecture.py", "projects/cosmos3/tokenizer/models/dense_backends.py": "cosmos_framework/model/tokenizer/models/dense_backends.py", "projects/cosmos3/tokenizer/models/dense_runtime.py": "cosmos_framework/model/tokenizer/models/dense_runtime.py", "projects/cosmos3/tokenizer/models/modules/__init__.py": "cosmos_framework/model/tokenizer/models/modules/__init__.py", @@ -445,6 +451,7 @@ "projects/cosmos3/tokenizer/models/text_decoder.py": "cosmos_framework/model/tokenizer/models/text_decoder.py", "projects/cosmos3/tokenizer/models/utils.py": "cosmos_framework/model/tokenizer/models/utils.py", "projects/cosmos3/tokenizer/utils/hf.py": "cosmos_framework/model/tokenizer/utils/hf.py", + "projects/cosmos3/tokenizer/utils/tensors.py": "cosmos_framework/model/tokenizer/utils/tensors.py", "projects/cosmos3/tokenizer/utils/vlm_prompt_format.py": "cosmos_framework/model/tokenizer/utils/vlm_prompt_format.py", "projects/cosmos3/utils/torchcodec_video.py": "cosmos_framework/utils/generator/torchcodec_video.py", "projects/cosmos3/vlm/configs/base/defaults/checkpointer.py": "cosmos_framework/utils/reasoner/configs_defaults/checkpointer.py", diff --git a/cosmos_framework/callbacks/every_n_draw_sample.py b/cosmos_framework/callbacks/every_n_draw_sample.py index fa87fda5..015c5f7f 100644 --- a/cosmos_framework/callbacks/every_n_draw_sample.py +++ b/cosmos_framework/callbacks/every_n_draw_sample.py @@ -217,6 +217,91 @@ def _add_wandb_image_paths( info[key_prefix] = wandb.Image(image_paths, caption=caption) +def _pixel_tensor_to_5d(t: torch.Tensor) -> torch.Tensor: + """Ensure a pixel tensor has shape (B, C, T, H, W) for the visualization grid. + + Handles (C, H, W), (B, C, H, W), and (B, C, T, H, W) inputs. + """ + if t.ndim == 3: + return t.unsqueeze(0).unsqueeze(2) # (C,H,W) -> (1,C,1,H,W) + if t.ndim == 4: + return t.unsqueeze(2) # (B,C,H,W) -> (B,C,1,H,W) + return t + + +def _resize_5d_to_width(img5d: torch.Tensor, target_width: int) -> torch.Tensor: + """Resize a single-frame (1, C, 1, H, W) tensor so its width is exactly ``target_width``. + + Height is scaled proportionally to preserve the overall aspect ratio. Assumes a + single temporal frame (T == 1). + """ + h, w = img5d.shape[-2], img5d.shape[-1] + if w == target_width: + return img5d + new_h = max(1, round(h * target_width / w)) + chw = img5d[0, :, 0] # (C,H,W) + resized = torchvision_F.resize(chw, [new_h, target_width], antialias=True) # (C,new_h,target_width) + return resized.unsqueeze(0).unsqueeze(2) # (1,C,1,new_h,target_width) + + +def _resize_pad_to_square(img5d: torch.Tensor, cell: int) -> torch.Tensor: + """Resize a single-frame (1, C, 1, H, W) tensor to fit inside a ``cell`` x ``cell`` square. + + Aspect ratio is preserved (the image is scaled so its longer side equals ``cell``), then + the result is center-padded with zeros to exactly ``cell`` x ``cell``. Assumes T == 1. + """ + h, w = img5d.shape[-2], img5d.shape[-1] + scale = cell / max(h, w) + new_h = max(1, round(h * scale)) + new_w = max(1, round(w * scale)) + chw = img5d[0, :, 0] # (C,H,W) + resized = torchvision_F.resize(chw, [new_h, new_w], antialias=True) # (C,new_h,new_w) + pad_h = cell - new_h + pad_w = cell - new_w + top = pad_h // 2 + left = pad_w // 2 + padded = torch.nn.functional.pad(resized, (left, pad_w - left, top, pad_h - top)) # (C,cell,cell) + return padded.unsqueeze(0).unsqueeze(2) # (1,C,1,cell,cell) + + +def _build_reference_grid(references: list[torch.Tensor], target_width: int) -> torch.Tensor: + """Tile reference images into a compact near-square grid sized to ``target_width``. + + All references are arranged in a ``rows`` x ``cols`` grid (``cols = ceil(sqrt(n))``), each in + an aspect-preserved square cell, then the whole grid is resized so its width equals + ``target_width``. This keeps the condition montage aligned with the generated-image column + width (better space utilization) instead of one overly-long row. + + Args: + references: list of single-frame pixel tensors (each (1,C,1,H,W) or (C,H,W)), in [-1, 1]. + target_width: width of the generated/target image for this sample. + + Returns: + A (1, C, 1, H_grid, target_width) tensor. + """ + if not references: + raise ValueError("Expected at least one reference image to build the condition grid.") + + refs = [_pixel_tensor_to_5d(r) for r in references] # list[(1,C,1,H,W)] + n = len(refs) + cols = math.ceil(math.sqrt(n)) + rows = math.ceil(n / cols) + cell = max(1, target_width // cols) + + cells = [_resize_pad_to_square(r, cell) for r in refs] # list[(1,C,1,cell,cell)] + blank = torch.zeros_like(cells[0]) + while len(cells) < rows * cols: + cells.append(blank) + + row_imgs = [] + for r in range(rows): + row = torch.cat(cells[r * cols : (r + 1) * cols], dim=-1) # (1,C,1,cell,cols*cell) + row_imgs.append(row) + grid = torch.cat(row_imgs, dim=-2) # (1,C,1,rows*cell,cols*cell) + + return _resize_5d_to_width(grid, target_width) # (1,C,1,H_grid,target_width) + + class EveryNDrawSample(EveryN): """ This callback sample condition inputs from training data, run inference and save the results to wandb and s3. @@ -531,11 +616,20 @@ def sample( vis_offset = 0 for sample_idx in range(data_clean.batch_size): n_vis = num_items[sample_idx] - # First item(s) are condition, last item is generation target - # but we need to support multiple conditions per sample in the future. Current code - # can handle this without throwing an error. - condition_images.append(raw_data[vis_offset]) # source image (1, C, 1, H, W) - gt_target_images.append(raw_data[vis_offset + n_vis - 1]) # target image (1, C, 1, H, W) + # First item(s) are condition references, last item is the generation target. + refs = raw_data[vis_offset : vis_offset + n_vis - 1] # all condition items + target = raw_data[vis_offset + n_vis - 1] # target image (1, C, 1, H, W) + # Multi-reference generation (>1 single-frame image references): tile every + # reference into a compact grid resized to the target width, so all references + # are visible without blowing up the row width. For video editing/transfer + # (T > 1) keep the existing behavior (first item = condition) so those tasks + # render exactly as before and stay consistent with the t_crop frame cropping. + refs_are_images = all(r.shape[-3] == 1 for r in refs) and target.shape[-3] == 1 + if refs_are_images and len(refs) > 1: + condition_images.append(_build_reference_grid(refs, target.shape[-1])) + else: + condition_images.append(raw_data[vis_offset]) # source image (1, C, 1, H, W) / video clip + gt_target_images.append(target) vis_offset += n_vis # Use target images for max_w/max_h/t_crop (generated samples match target size) diff --git a/cosmos_framework/checkpoint/dcp.py b/cosmos_framework/checkpoint/dcp.py index 54d0d8e1..43145c88 100644 --- a/cosmos_framework/checkpoint/dcp.py +++ b/cosmos_framework/checkpoint/dcp.py @@ -71,9 +71,9 @@ from cosmos_framework.checkpoint.base import AbstractCheckpointer from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader, S3StorageWriter +from cosmos_framework.utils.config import CheckpointConfig, JobConfig from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.utils import callback, distributed, log, misc -from cosmos_framework.utils.config import CheckpointConfig, JobConfig from cosmos_framework.utils.easy_io import easy_io from cosmos_framework.utils.generator.rand_state import get_rand_state_dict, set_rand_state_dict @@ -866,16 +866,6 @@ def load( raise ValueError( f"Unexpected keys (found in checkpoint but not in model): {results.unexpected_keys}" ) - # Warm start that skipped net_ema (e.g. loading an EMA-only HF export - # with no net_ema.* keys): the EMA shadow would otherwise keep its random - # build-time generation pathway (init_moe is skipped when a checkpoint is - # present). Seed net_ema from the freshly loaded net so the EMA starts equal - # to net ("EMA warm-starts from net") instead of from random weights. - if warm_start and any("net_ema" in skip_key for skip_key in keys_to_skip_loading): - ema_worker = getattr(model, "net_ema_worker", None) - if ema_worker is not None and getattr(model, "net_ema", None) is not None: - ema_worker.copy_to(src_model=model.net, tgt_model=model.net_ema) - log.info("Warm start: re-seeded net_ema from net (net_ema was skipped on load).") elif key == "optim": log.info("- Loading the optimizer...") diff --git a/cosmos_framework/configs/base/defaults/checkpointer.py b/cosmos_framework/configs/base/defaults/checkpointer.py index 9502070e..151f6e61 100644 --- a/cosmos_framework/configs/base/defaults/checkpointer.py +++ b/cosmos_framework/configs/base/defaults/checkpointer.py @@ -98,6 +98,7 @@ strict_resume=True, enable_gcs_patch_in_boto3=True, dcp_async_mode_enabled=True, + dcp_load_dedup=True, ) CHECKPOINT_S3_EAST2 = CheckpointConfig( diff --git a/cosmos_framework/configs/base/defaults/model_config.py b/cosmos_framework/configs/base/defaults/model_config.py index adde6bd9..5e6fd6b4 100644 --- a/cosmos_framework/configs/base/defaults/model_config.py +++ b/cosmos_framework/configs/base/defaults/model_config.py @@ -5,12 +5,12 @@ import attrs +from cosmos_framework.utils.lazy_config import LazyDict from cosmos_framework.configs.base.defaults.activation_checkpointing import ActivationCheckpointingConfig from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.ema import EMAConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig from cosmos_framework.configs.base.defaults.reasoner import VLMConfig -from cosmos_framework.utils.lazy_config import LazyDict @attrs.define(slots=False) diff --git a/cosmos_framework/configs/base/defaults/optimizer.py b/cosmos_framework/configs/base/defaults/optimizer.py index dc56906c..a619b2fe 100644 --- a/cosmos_framework/configs/base/defaults/optimizer.py +++ b/cosmos_framework/configs/base/defaults/optimizer.py @@ -31,6 +31,57 @@ disable_weight_decay_for_1d_params=False, ) +# Muon / Dion2 share the standard factory knobs (keys_to_select, lr_multipliers, +# disable_weight_decay_for_1d_params) plus their own orthogonalization +# hyperparameters. ``fused`` is required by the factory; the AdamW side is fused +# by construction and ``capturable`` / ``master_weights`` are forced on. +MUON_OPTIMIZER_KWARGS: dict[str, Any] = dict( + # Base learning rate. Muon scales matrix params by muon_lr_scale*sqrt(max(A,B)); + # the AdamW side and the per-param-group lr_multipliers use it directly. + lr=1e-4, + weight_decay=0.1, + adam_betas=[0.9, 0.99], + eps=1e-8, + fused=True, + keys_to_select=[], + lr_multipliers={}, + disable_weight_decay_for_1d_params=False, + # Name substrings for stacked MoE expert params ([E, M, N]) to orthogonalize + # per expert slice. Empty = experts stay on AdamW (no behavior change). + # e.g. ["gate_up_proj", "down_proj"] for grouped-MM MoE experts. + expert_param_keywords=[], + # Muon-specific. + muon_momentum=0.95, + muon_lr_scale=0.2, + ns_steps=5, + nesterov=True, + use_distributed=True, +) + +DION2_OPTIMIZER_KWARGS: dict[str, Any] = dict( + lr=1e-4, + weight_decay=0.1, + adam_betas=[0.9, 0.99], + eps=1e-8, + fused=True, + keys_to_select=[], + lr_multipliers={}, + disable_weight_decay_for_1d_params=False, + # Name substrings for stacked MoE expert params ([E, M, N]) to orthogonalize + # per expert slice. Empty = experts stay on AdamW (no behavior change). + # e.g. ["gate_up_proj", "down_proj"] for grouped-MM MoE experts. + expert_param_keywords=[], + # Muon/Dion2-specific. + muon_momentum=0.95, + muon_lr_scale=0.2, + ns_steps=5, + nesterov=True, + use_distributed=True, + # Dion2-specific: submatrix selection fraction and error-feedback decay. + fraction=1.0, + ef_decay=0.95, +) + LAMBDACOSINE_KWARGS: dict[str, Any] = dict( warm_up_steps=[2000], cycle_lengths=[100000], @@ -64,6 +115,26 @@ def register_optimizers(optimizer_kwargs: dict[str, Any]) -> None: **optimizer_kwargs, ), ) + cs.store( + group="optimizer", + package="optimizer", + name="muonwithauxadamw", + node=L(build_optimizer)( + model=PLACEHOLDER, + optimizer_type="MuonWithAuxAdamW", + **MUON_OPTIMIZER_KWARGS, + ), + ) + cs.store( + group="optimizer", + package="optimizer", + name="dion2withauxadamw", + node=L(build_optimizer)( + model=PLACEHOLDER, + optimizer_type="Dion2WithAuxAdamW", + **DION2_OPTIMIZER_KWARGS, + ), + ) def register_schedulers(lambdacosine_kwargs: dict[str, Any]) -> None: diff --git a/cosmos_framework/configs/base/defaults/reasoner.py b/cosmos_framework/configs/base/defaults/reasoner.py index bee81b63..09e92efe 100644 --- a/cosmos_framework/configs/base/defaults/reasoner.py +++ b/cosmos_framework/configs/base/defaults/reasoner.py @@ -785,6 +785,32 @@ class VLMConfig: ), ) +# Same as 590c1c0 but with use_und_k_norm_for_gen=True: normalises K_und before +# it is used as a key in the gen→und cross-attention path (the qk-norm fix for the +# generator). Adds a freshly-initialised k_norm_und_for_gen RMSNorm. +Cosmos3EdgeReasoner_VLM_GCP_Config_590c1c0_UndKNorm: VLMConfig = VLMConfig( + model_name="nvidia/Cosmos3-Edge-Reasoner", + model_instance=L(Nemotron3DenseVLTextForCausalLM)( + config=L(create_vlm_config)( + base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( + json_file="cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + ), + qk_norm_for_text=False, + use_und_k_norm_for_gen=True, + ), + ), + tokenizer=L(build_processor_lazy)( + tokenizer_type="nvidia/Cosmos3-Edge-Reasoner", + config_variant="gcp", + ), + pretrained_weights=PretrainedWeightsConfig( + backbone_path="s3://bucket0/cosmos3/pretrained/huggingface/nvidia/Cosmos3-Edge-Reasoner-590c1c0/", + credentials_path="credentials/gcp_checkpoint.secret", + enable_gcs_patch_in_boto3=True, + checkpoint_format="nemotron_3_dense_vl", + ), +) + # Same as 9b4c028 but with use_und_k_norm_for_gen=True: normalises K_und before # it is used as a key in the gen→und cross-attention path. Cosmos3EdgeReasoner_VLM_GCP_Config_9b4c028_UndKNorm: VLMConfig = VLMConfig( @@ -987,3 +1013,9 @@ def register_vlm(): name="cosmos3_edge_reasoner_vlm_gcp_590c1c0", node=Cosmos3EdgeReasoner_VLM_GCP_Config_590c1c0, ) + cs.store( + group="vlm_config", + package="model.config.vlm_config", + name="cosmos3_edge_reasoner_vlm_gcp_590c1c0_und_k_norm", + node=Cosmos3EdgeReasoner_VLM_GCP_Config_590c1c0_UndKNorm, + ) diff --git a/cosmos_framework/configs/base/defaults/tokenizer.py b/cosmos_framework/configs/base/defaults/tokenizer.py index 04c08873..55e79016 100644 --- a/cosmos_framework/configs/base/defaults/tokenizer.py +++ b/cosmos_framework/configs/base/defaults/tokenizer.py @@ -128,6 +128,7 @@ vae_path=PRETRAINED_TOKENIZER_UNIAE_4X16X16_C48_T16TO160_MIXP_FPS_MIX_ENCODER_NONCAUSAL_DECODER_NONCAUSAL_NOGAN_S3_NEMOTRON2B_VAE_PTH, spatial_compression_factor=16, temporal_compression_factor=4, + pad_frames=1, pixel_trim=True, causal=False, ) diff --git a/cosmos_framework/data/generator/action/domain_utils.py b/cosmos_framework/data/generator/action/domain_utils.py index 6f433f73..c75f7f38 100644 --- a/cosmos_framework/data/generator/action/domain_utils.py +++ b/cosmos_framework/data/generator/action/domain_utils.py @@ -20,6 +20,7 @@ "agibotworld": 15, "embodiment_c_gripper": 15, "embodiment_c_gripper_ext": 15, + "xdof_yam": 16, "fractal": 20, } @@ -38,6 +39,7 @@ "agibotworld": 29, "embodiment_c_gripper": 29, "embodiment_c_gripper_ext": 29, + "xdof_yam": 20, "fractal": 10, # NOTE: ``libero`` (7/10/13 depending on ``rotation_space``) and ``hand_pose`` # (variable with ``keypoint_option`` and ``rotation_format``) are absent diff --git a/cosmos_framework/data/generator/utils.py b/cosmos_framework/data/generator/utils.py index bba5b3de..ecfcc0fe 100644 --- a/cosmos_framework/data/generator/utils.py +++ b/cosmos_framework/data/generator/utils.py @@ -98,7 +98,10 @@ def parse_frame_range_from_wdinfo(wdinfo: str) -> tuple[int, int | float] | None wdinfo: wdinfo path string containing a frames_X_Y pattern, where Y may be ``inf`` Returns: - Tuple of (min_frames, max_frames) if found, None otherwise + Tuple of (min_frames, max_frames) if found, where max_frames is + ``math.inf`` for unbounded ``frames_X_inf`` and ``frames_gt_X`` buckets. + ``frames_gt_X`` is treated as starting at X. Returns None if no frame + range is present. Example: >>> parse_frame_range_from_wdinfo("wdinfo/v4/tv_drama/resolution_720/aspect_ratio_16_9/frames_300_400/wdinfo.json") @@ -110,6 +113,10 @@ def parse_frame_range_from_wdinfo(wdinfo: str) -> tuple[int, int | float] | None if match: max_frames = math.inf if match.group(2) == "inf" else int(match.group(2)) return (int(match.group(1)), max_frames) + match = re.search(r"frames_gt_(\d+)", wdinfo) + if match: + # Wdinfo generation uses frames_gt_X for the bucket that starts at X, so treat it as >= X. + return (int(match.group(1)), math.inf) return None diff --git a/cosmos_framework/model/attention/backends.py b/cosmos_framework/model/attention/backends.py index 4236b35b..8d5f83fb 100644 --- a/cosmos_framework/model/attention/backends.py +++ b/cosmos_framework/model/attention/backends.py @@ -10,6 +10,7 @@ import torch +from cosmos_framework.model.attention.cudnn.checks import cudnn_attention_check from cosmos_framework.model.attention.flash2.checks import flash2_attention_check from cosmos_framework.model.attention.flash3.checks import flash3_attention_check from cosmos_framework.model.attention.masks import CausalType @@ -22,8 +23,8 @@ from cosmos_framework.model.attention.utils.safe_ops import log from cosmos_framework.model.attention.utils.safe_ops.functools import lru_cache - BACKEND_CHECK_MAP = { + "cudnn": cudnn_attention_check, "natten": natten_attention_check, "flash2": flash2_attention_check, "flash3": flash3_attention_check, @@ -131,17 +132,25 @@ def get_backend_list(arch_tag: int) -> list[str]: if arch_tag == 90: default_backends = [ "flash3", + "cudnn", "natten", "flash2", ] elif arch_tag in [100, 103]: default_backends = [ + "cudnn", "natten", "flash2", ] + elif arch_tag in [110, 120, 121]: + default_backends = [ + "cudnn", + "natten", + ] elif arch_tag >= 80: default_backends = [ "flash2", + "cudnn", "natten", ] else: diff --git a/cosmos_framework/model/attention/benchmarks/benchmark_fmha.py b/cosmos_framework/model/attention/benchmarks/benchmark_fmha.py new file mode 100644 index 00000000..bc68b1c5 --- /dev/null +++ b/cosmos_framework/model/attention/benchmarks/benchmark_fmha.py @@ -0,0 +1,638 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +""" +Benchmark the unified ``cosmos_framework.model.attention.attention`` frontend across a +user-specified list of backends, for a single sequence-length (and optional +context-parallel shard shape). + +Dense / batched attention ONLY: this benchmark deliberately has no varlen +(sequence-packed) support. Inputs are ``[B, S, H, D]``. + +This single script combines two previously-separate benchmarks: + +1. Backend sweep (the former ``benchmark_fmha``): compare backends at a single query + length (``--seqlen``), with optional cross-attention (``--kv-seqlen``), causal + masking, backward-pass timing, MLA-style value head dim (``--head-dim-v``), and + per-backend rejection diagnostics (``--debug``). + +2. Context-parallel per-rank shape study : + isolate the per-rank attention kernel shape a CP strategy would produce, via + ``--shard-mode`` + ``--cp-size``. ``head`` divides Q/KV heads by ``cp_size`` and + keeps the full KV length; ``sequence`` keeps the heads and splits the KV length + across ``cp_size`` shards (Q is intentionally NOT divided — each rank computes a + partial result over a local KV shard). There is intentionally no all-to-all, + all-gather, or all-reduce in the timed region: the intent is to study local kernel + scalability. The shard transform is applied to the ``(q_len, kv_len)`` shape. With + ``--shard-mode none`` (default) no sharding is applied. + +The script runs single-process by default, but also runs under ``torchrun`` so a real +multi-rank CP shape can be measured (``sequence`` mode is rank-dependent). Metrics +reported per row: median / min latency, achieved TFLOP/s (on the LOCAL shape), tokens/s, +and peak allocated memory. Only rank 0 prints. + +By default this is self-attention (``q_len == kv_len``). Cross-attention with a +different KV length can be benchmarked via ``--kv-seqlen``, but only when ``--causal`` +is NOT set: causal masking here uses ``CausalType.DontCare``, which requires +``q_len == kv_len``. When ``--causal`` is set, ``kv_len`` always equals ``q_len``. + +Examples: + + # Backend sweep, self-attention, causal, forward-only: + python -m cosmos_framework.model.attention.benchmarks.benchmark_fmha \ + --backends natten cudnn \ + --seqlen 3093 \ + --batch 1 --num-q-heads 16 --num-kv-heads 8 --head-dim 128 \ + --dtype bf16 --causal + + # Cross-attention (non-causal), distinct q_len and kv_len: + python -m cosmos_framework.model.attention.benchmarks.benchmark_fmha \ + --backends natten cudnn \ + --seqlen 4096 --kv-seqlen 1024 \ + --num-q-heads 16 --num-kv-heads 16 --head-dim 128 --dtype bf16 + + # One-GPU mock of a CP=4 per-rank head-sharded shape: + python -m cosmos_framework.model.attention.benchmarks.benchmark_fmha \ + --backends natten cudnn \ + --seqlen 396 --kv-seqlen 177771 \ + --num-q-heads 32 --num-kv-heads 8 --head-dim 128 \ + --shard-mode head --cp-size 4 --compile + + # True four-GPU CP head-sharded run: + torchrun --standalone --nproc_per_node=4 \ + -m cosmos_framework.model.attention.benchmarks.benchmark_fmha \ + --seqlen 396 --kv-seqlen 177771 \ + --num-q-heads 32 --num-kv-heads 8 --head-dim 128 \ + --shard-mode head --cp-size 4 + + # One-GPU mock of a CP=4 per-rank sequence-sharded (split-KV) shape. + # KV length is divided by cp_size while heads are kept; --causal is not allowed: + python -m cosmos_framework.model.attention.benchmarks.benchmark_fmha \ + --backends natten cudnn \ + --seqlen 396 --kv-seqlen 177771 \ + --num-q-heads 32 --num-kv-heads 8 --head-dim 128 \ + --shard-mode sequence --cp-size 4 + + # True four-GPU CP sequence-sharded run (each rank holds a different KV shard): + torchrun --standalone --nproc_per_node=4 \ + -m cosmos_framework.model.attention.benchmarks.benchmark_fmha \ + --seqlen 396 --kv-seqlen 177771 \ + --num-q-heads 32 --num-kv-heads 8 --head-dim 128 \ + --shard-mode sequence --cp-size 4 + +Pass ``--debug`` to report the concrete reason an incompatible backend was rejected +(rather than the generic "incompatible" message) and to print DEBUG-level +backend-selection logs. Pass ``--json`` to emit one JSON object per result row instead +of the formatted table. +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch import Tensor + +from cosmos_framework.model.attention import attention +from cosmos_framework.model.attention.backends import choose_backend +from cosmos_framework.model.attention.masks import CausalType +from cosmos_framework.model.attention.utils import get_arch_tag + +DTYPE_MAP: dict[str, torch.dtype] = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, +} + + +@dataclass(frozen=True) +class LocalAttentionShape: + """Per-rank attention shape after applying the CP shard transform to a global pair.""" + + q_len: int + kv_len: int + num_q_heads: int + num_kv_heads: int + sequence_shard_index: int + sequence_shard_start: int + + +def local_shape_for_pair( + q_len: int, + kv_len: int, + num_q_heads: int, + num_kv_heads: int, + shard_mode: str, + cp_size: int, + rank: int, +) -> LocalAttentionShape: + """Transform a global ``(q_len, kv_len, heads)`` pair into the LOCAL per-rank shape. + + - ``none``: identity (cp_size must be 1). + - ``head``: divide Q and KV heads by ``cp_size``, keep the full KV length. All ranks + see the same local shape (only head counts shrink). + - ``sequence``: keep the heads, split the KV length across ``cp_size`` shards; the + local KV length depends on ``rank % cp_size`` (remainder distributed to the low + shards). Q is intentionally not divided. + """ + if shard_mode == "none": + if cp_size != 1: + raise ValueError(f"--shard-mode none requires --cp-size 1, got {cp_size}") + return LocalAttentionShape(q_len, kv_len, num_q_heads, num_kv_heads, 0, 0) + + if shard_mode == "head": + if num_q_heads % cp_size != 0: + raise ValueError(f"num_q_heads={num_q_heads} must be divisible by cp_size={cp_size}") + if num_kv_heads % cp_size != 0: + raise ValueError(f"num_kv_heads={num_kv_heads} must be divisible by cp_size={cp_size}") + local_q_heads = num_q_heads // cp_size + local_kv_heads = num_kv_heads // cp_size + if local_q_heads % local_kv_heads != 0: + raise ValueError(f"local_q_heads={local_q_heads} must be divisible by local_kv_heads={local_kv_heads}") + return LocalAttentionShape(q_len, kv_len, local_q_heads, local_kv_heads, 0, 0) + + if shard_mode == "sequence": + shard_index = rank % cp_size + base_kv_len = kv_len // cp_size + remainder = kv_len % cp_size + local_kv_len = base_kv_len + int(shard_index < remainder) + if local_kv_len <= 0: + raise ValueError( + f"sequence-sharded local_kv_len must be positive, got {local_kv_len}; " + f"kv_len={kv_len}, cp_size={cp_size}, shard_index={shard_index}" + ) + sequence_shard_start = shard_index * base_kv_len + min(shard_index, remainder) + return LocalAttentionShape(q_len, local_kv_len, num_q_heads, num_kv_heads, shard_index, sequence_shard_start) + + raise ValueError(f"Unsupported shard_mode={shard_mode!r}") + + +def attention_flops( + batch: int, + q_len: int, + kv_len: int, + num_q_heads: int, + head_dim: int, + head_dim_v: int, + is_causal: bool, + include_backward: bool, +) -> float: + """ + Approximate attention FLOPs (supports q_len != kv_len for cross-attention). + + Forward is 2 GEMMs (QK^T and P@V); each multiply-add counts as 2 FLOPs. + Backward is ~2x the forward cost. Causal masking roughly halves the work + (only applicable when q_len == kv_len). + """ + qk = 2.0 * batch * num_q_heads * q_len * kv_len * head_dim # [B,Hq,Sq,Skv] scores + pv = 2.0 * batch * num_q_heads * q_len * kv_len * head_dim_v # [B,Hq,Sq,Dv] output + flops = qk + pv + if is_causal: + flops *= 0.5 + if include_backward: + flops *= 3.0 # fwd (1x) + bwd (~2x) + return flops + + +def make_inputs( + batch: int, + q_len: int, + kv_len: int, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + head_dim_v: int, + dtype: torch.dtype, + device: torch.device, + requires_grad: bool, + generator: torch.Generator | None, +) -> tuple[Tensor, Tensor, Tensor]: + """Create heads-last dense QKV tensors (q_len may differ from kv_len).""" + q = torch.randn( + batch, q_len, num_q_heads, head_dim, dtype=dtype, device=device, generator=generator + ).requires_grad_(requires_grad) # [B,Sq,Hq,D] + k = torch.randn( + batch, kv_len, num_kv_heads, head_dim, dtype=dtype, device=device, generator=generator + ).requires_grad_(requires_grad) # [B,Skv,Hkv,D] + v = torch.randn( + batch, kv_len, num_kv_heads, head_dim_v, dtype=dtype, device=device, generator=generator + ).requires_grad_(requires_grad) # [B,Skv,Hkv,Dv] + return q, k, v + + +def compatibility_reason( + q: Tensor, # [B,S,Hq,D] + k: Tensor, # [B,S,Hkv,D] + v: Tensor, # [B,S,Hkv,Dv] + backend: str | None, + is_causal: bool, +) -> str | None: + """ + Surface the concrete reason a backend was rejected for this use case / device. + + The frontend runs backend selection with ``raise_error=False`` and only reports a + generic "incompatible" message, swallowing the underlying reason. Here we re-run + ``choose_backend`` with ``raise_error=True`` (a distinct lru_cache key, so it does not + hit the frontend's cached result) to capture the specific rejection reason. + + Returns the reason string, or None if the backend is compatible. + """ + causal_type = CausalType.DontCare if is_causal else None + try: + choose_backend( + query_shape=q.shape, + key_shape=k.shape, + value_shape=v.shape, + dtype=q.dtype, + device=q.device, + requires_grad=q.requires_grad or k.requires_grad or v.requires_grad, + is_causal=is_causal, + causal_type=causal_type, + is_varlen=False, + deterministic=False, + backend=backend, + raise_error=True, + ) + return None + except Exception as e: # noqa: BLE001 - surface the concrete rejection reason + return f"{type(e).__name__}: {e}" + + +def run_once( + q: Tensor, # [B,S,Hq,D] + k: Tensor, # [B,S,Hkv,D] + v: Tensor, # [B,S,Hkv,Dv] + backend: str | None, + is_causal: bool, + include_backward: bool, +) -> Tensor: + """Run a single forward (and optional backward) attention call. + + Returns the attention output. The caller MUST keep/consume the returned tensor: + under ``torch.compile`` an unused output lets Inductor dead-code-eliminate the whole + attention (yielding absurd, ~PFLOP/s "timings"), so the output is returned as a graph + output to keep the kernel live. + """ + causal_type = CausalType.DontCare if is_causal else None + out = attention( + query=q, + key=k, + value=v, + is_causal=is_causal, + causal_type=causal_type, + backend=backend, + return_lse=False, + ) # [B,S,Hq,Dv] + if include_backward: + grad = torch.ones_like(out) # [B,S,Hq,Dv] + out.backward(grad) + return out + + +def benchmark_backend( + q: Tensor, # [B,S,Hq,D] + k: Tensor, # [B,S,Hkv,D] + v: Tensor, # [B,S,Hkv,Dv] + backend: str | None, + is_causal: bool, + include_backward: bool, + warmup: int, + iters: int, + compile_fn: bool, + shard_mode: str, + cp_size: int, + debug: bool, +) -> tuple[list[float], int, str | None]: + """ + Time a single backend on a single (already-sharded) input shape. + + Returns ``(per-iter latencies in ms, peak allocated bytes, error message)``. + If the backend is incompatible or errors out, latencies is empty, peak bytes is 0, + and the error message describes why. When ``debug`` is set, an incompatible backend's + concrete rejection reason is reported instead of the generic frontend message. + + The forward (+ optional backward) is wrapped in an NVTX range so the timed region is + easy to find in Nsight Systems. Backward requires grad, so ``inference_mode`` is only + used for the forward-only path. ``compile_fn`` compiles the call before timing. + """ + if debug: + reason = compatibility_reason(q, k, v, backend, is_causal) + if reason is not None: + return [], 0, reason + + def _call() -> Tensor: + return run_once(q, k, v, backend, is_causal, include_backward) + + call = torch.compile(_call, fullgraph=not include_backward) if compile_fn else _call + + # inference_mode is incompatible with autograd; only use it for forward-only timing. + ctx = torch.inference_mode() if not include_backward else torch.enable_grad() + + if dist.is_initialized(): + dist.barrier() + + # Sink holding the latest output so Inductor cannot dead-code-eliminate the attention + # (an unused compiled output would be optimized away, giving impossible TFLOP/s). + sink: Tensor | None = None + try: + with ctx: + for _ in range(warmup): + sink = call() + torch.cuda.synchronize() + + torch.cuda.reset_peak_memory_stats(q.device) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + latencies_ms: list[float] = [] + torch.cuda.nvtx.range_push(f"fmha.{shard_mode}.cp{cp_size}.{backend or 'auto'}") + for _ in range(iters): + start.record() + sink = call() + end.record() + torch.cuda.synchronize() + latencies_ms.append(start.elapsed_time(end)) # ms + torch.cuda.nvtx.range_pop() + # Force a host-side read of the result so the compiled output is genuinely + # materialized (belt-and-suspenders against dead-code elimination). + if sink is not None: + _ = sink.sum().item() + except Exception as e: # noqa: BLE001 - report and skip incompatible backends + return [], 0, f"{type(e).__name__}: {e}" + + peak_bytes = torch.cuda.max_memory_allocated(q.device) + return latencies_ms, peak_bytes, None + + +def format_row(cells: list[str], widths: list[int]) -> str: + return " ".join(cell.ljust(w) for cell, w in zip(cells, widths)) + + +def init_distributed() -> tuple[int, int, int]: + """Read torchrun env (defaults to single process) and init NCCL when world_size > 1.""" + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + if world_size > 1 and not dist.is_initialized(): + dist.init_process_group("nccl") + return rank, local_rank, world_size + + +def resolve_seqlen_pair(seqlen: int, kv_seqlen: int | None, causal: bool) -> tuple[int, int]: + """Resolve the single (q_len, kv_len) shape. Causal masking here uses + CausalType.DontCare, which requires q_len == kv_len, so a distinct KV length is only + allowed when not causal.""" + if kv_seqlen is None: + return (seqlen, seqlen) + if causal: + raise SystemExit( + "--kv-seqlen is only valid without --causal: causal masking here uses CausalType.DontCare, " + "which requires q_len == kv_len. Drop --kv-seqlen, or drop --causal." + ) + return (seqlen, kv_seqlen) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Benchmark cosmos_framework.model.attention (dense / non-varlen only) across backends, " + "sequence lengths, and optional context-parallel shard shapes." + ) + parser.add_argument( + "--backends", + nargs="+", + default=["auto"], + help="Backends to benchmark. Use 'auto' to let the frontend choose. " + "Named choices: flash2, flash3, natten, cudnn.", + ) + parser.add_argument("--seqlen", type=int, default=4096, help="Query sequence length.") + parser.add_argument( + "--kv-seqlen", + type=int, + default=None, + help="KV sequence length for cross-attention (q_len != kv_len). Only allowed when " + "--causal is NOT set. If omitted, kv_len == q_len (self-attention).", + ) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--num-q-heads", type=int, default=16) + parser.add_argument("--num-kv-heads", type=int, default=8) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument( + "--head-dim-v", type=int, default=None, help="Value head dim (defaults to --head-dim; set differently for MLA)." + ) + parser.add_argument("--dtype", choices=list(DTYPE_MAP), default="bf16") + parser.add_argument("--causal", action="store_true", help="Enable causal masking (CausalType.DontCare).") + parser.add_argument("--backward", action="store_true", help="Include a backward pass (training) in timing.") + parser.add_argument( + "--shard-mode", + choices=("none", "head", "sequence"), + default="none", + help="Context-parallel per-rank shape to simulate. none: no sharding; " + "head: divide Q/KV heads by --cp-size (full KV length); " + "sequence: keep heads and split KV length across --cp-size shards.", + ) + parser.add_argument( + "--cp-size", type=int, default=1, help="Context-parallel sharding factor (>1 needs --shard-mode)." + ) + parser.add_argument("--compile", action="store_true", help="torch.compile the attention call before benchmarking.") + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--seed", type=int, default=1234, help="Random seed for input generation.") + parser.add_argument("--device", default="cuda") + parser.add_argument("--json", action="store_true", help="Emit one JSON object per result row instead of a table.") + parser.add_argument( + "--debug", + action="store_true", + help="Report the concrete reason a backend is rejected (instead of the generic " + "'incompatible' message) and enable DEBUG-level backend-selection logs.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + if args.debug: + # Raise the logger to DEBUG so the per-backend selection diagnostics (emitted by the + # attention checks via log.debug) are printed, then re-initialize the stdout sink. + from cosmos_framework.utils import log + + log.LEVEL = "DEBUG" + log.init_loguru_stdout() + + if not torch.cuda.is_available(): + raise SystemExit("CUDA is not available; this benchmark requires a GPU.") + + if args.cp_size < 1: + raise SystemExit(f"--cp-size must be >= 1, got {args.cp_size}") + if args.shard_mode == "sequence" and args.causal: + raise SystemExit( + "--shard-mode sequence is incompatible with --causal: a KV-split shard holds a " + "partial key range, so a causal mask over the local shard is not well-defined here." + ) + + rank, local_rank, world_size = init_distributed() + is_main = rank == 0 + device = torch.device("cuda", local_rank) if args.device == "cuda" else torch.device(args.device) + dtype = DTYPE_MAP[args.dtype] + head_dim_v = args.head_dim_v if args.head_dim_v is not None else args.head_dim + backends: list[str | None] = [None if b == "auto" else b for b in args.backends] + generator = torch.Generator(device=device).manual_seed(args.seed) + + q_len, kv_len = resolve_seqlen_pair(args.seqlen, args.kv_seqlen, args.causal) + local = local_shape_for_pair( + q_len=q_len, + kv_len=kv_len, + num_q_heads=args.num_q_heads, + num_kv_heads=args.num_kv_heads, + shard_mode=args.shard_mode, + cp_size=args.cp_size, + rank=rank, + ) + seq_label = str(local.q_len) if local.q_len == local.kv_len else f"{local.q_len}/{local.kv_len}" + heads_label = f"{local.num_q_heads}/{local.num_kv_heads}" + shard_label = "-" if args.shard_mode == "none" else f"{args.shard_mode}/{args.cp_size}" + + if is_main and not args.json: + arch_tag = get_arch_tag(device) + print(f"Device: {torch.cuda.get_device_name(device)} (arch_tag={arch_tag}, sm_{arch_tag})") + print( + f"Config: batch={args.batch} Hq={args.num_q_heads} Hkv={args.num_kv_heads} " + f"D={args.head_dim} Dv={head_dim_v} dtype={args.dtype} causal={args.causal} " + f"backward={args.backward} compile={args.compile}" + ) + print( + f"Shard: mode={args.shard_mode} cp_size={args.cp_size} world_size={world_size} | " + f"Timing: warmup={args.warmup} iters={args.iters}\n" + ) + headers = [ + "backend", + "seq(q/kv)", + "heads(q/kv)", + "shard", + "median(ms)", + "min(ms)", + "TFLOP/s", + "tok/s", + "peakMB", + "status", + ] + widths = [10, 15, 12, 10, 11, 10, 9, 12, 9, 40] + print(format_row(headers, widths)) + print(format_row(["-" * w for w in widths], widths)) + else: + widths = [10, 15, 12, 10, 11, 10, 9, 12, 9, 40] + + for backend in backends: + backend_label = backend if backend is not None else "auto" + + q, k, v = make_inputs( + batch=args.batch, + q_len=local.q_len, + kv_len=local.kv_len, + num_q_heads=local.num_q_heads, + num_kv_heads=local.num_kv_heads, + head_dim=args.head_dim, + head_dim_v=head_dim_v, + dtype=dtype, + device=device, + requires_grad=args.backward, + generator=generator, + ) + + latencies_ms, peak_bytes, error = benchmark_backend( + q, + k, + v, + backend=backend, + is_causal=args.causal, + include_backward=args.backward, + warmup=args.warmup, + iters=args.iters, + compile_fn=args.compile, + shard_mode=args.shard_mode, + cp_size=args.cp_size, + debug=args.debug, + ) + + if not is_main: + continue + + if error is not None: + if args.json: + print(json.dumps({"backend": backend_label, "seq": seq_label, "status": error})) + else: + print( + format_row( + [backend_label, seq_label, heads_label, shard_label, "-", "-", "-", "-", "-", error], widths + ) + ) + continue + + median_ms = statistics.median(latencies_ms) + min_ms = min(latencies_ms) + flops = attention_flops( + batch=args.batch, + q_len=local.q_len, + kv_len=local.kv_len, + num_q_heads=local.num_q_heads, + head_dim=args.head_dim, + head_dim_v=head_dim_v, + is_causal=args.causal, + include_backward=args.backward, + ) + tflops = flops / (median_ms * 1e-3) / 1e12 + tokens_per_s = (args.batch * local.q_len) / (median_ms * 1e-3) + peak_mb = peak_bytes / (1024**2) + + if args.json: + print( + json.dumps( + { + "backend": backend_label, + "q_len": local.q_len, + "kv_len": local.kv_len, + "num_q_heads": local.num_q_heads, + "num_kv_heads": local.num_kv_heads, + "shard_mode": args.shard_mode, + "cp_size": args.cp_size, + "median_ms": median_ms, + "min_ms": min_ms, + "tflops": tflops, + "tokens_per_s": tokens_per_s, + "peak_mb": peak_mb, + "status": "ok", + } + ) + ) + else: + print( + format_row( + [ + backend_label, + seq_label, + heads_label, + shard_label, + f"{median_ms:.4f}", + f"{min_ms:.4f}", + f"{tflops:.1f}", + f"{tokens_per_s:.0f}", + f"{peak_mb:.1f}", + "ok", + ], + widths, + ) + ) + + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/cosmos_framework/model/attention/cudnn/__init__.py b/cosmos_framework/model/attention/cudnn/__init__.py index ca3697b9..c585d320 100644 --- a/cosmos_framework/model/attention/cudnn/__init__.py +++ b/cosmos_framework/model/attention/cudnn/__init__.py @@ -11,12 +11,10 @@ import torch from cosmos_framework.model.attention.utils.safe_ops import log -from cosmos_framework.model.attention.utils.version import version_at_least -CUDNN_DISALLOWED = True - -CUDNN_MIN_BACKEND_VERSION = 91300 -CUDNN_MIN_FRONTEND_VERSION = "1.14.0" +# Minimum cuDNN runtime version, in ``torch.backends.cudnn.version()`` encoding +# (major * 10000 + minor * 100 + patch). 92200 == cuDNN 9.22.0. +CUDNN_MIN_BACKEND_VERSION = 92200 def cudnn_supported() -> bool: @@ -24,36 +22,34 @@ def cudnn_supported() -> bool: Returns whether cuDNN Attention is supported in this environment. Requirements are: * Presence of CUDA Runtime (via PyTorch) - * Presence of cuDNN and its Python frontend, meeting minimum version requirements + * Presence of the cuDNN runtime that ships with / is linked by PyTorch, meeting the minimum + version requirement - This check guards imports / dependencies on the cuDNN package. + The backend runs cuDNN attention through PyTorch's own SDPA cuDNN dispatch + (``F.scaled_dot_product_attention`` / ``aten._scaled_dot_product_cudnn_attention``), so it no + longer depends on the standalone cuDNN Python frontend package -- only on the cuDNN that PyTorch + itself uses. """ if not torch.cuda.is_available(): log.debug("cuDNN Attention is not supported because PyTorch did not detect CUDA runtime.") return False - try: - import cudnn - - except ImportError: - log.debug("cuDNN Attention is not supported because the frontend Python package was not found.") - return False - except Exception as e: - log.debug(f"cuDNN Attention is not supported because importing the frontend Python package failed: {e}") + if not torch.backends.cudnn.is_available(): + log.debug("cuDNN Attention is not supported because PyTorch reports cuDNN is unavailable.") return False - if cudnn.backend_version() < CUDNN_MIN_BACKEND_VERSION: + backend_version = torch.backends.cudnn.version() + if backend_version is None or backend_version < CUDNN_MIN_BACKEND_VERSION: log.debug( - "cuDNN Attention is not supported due to insufficient cuDNN backend version " - f"{cudnn.backend_version()=}, expected at least {CUDNN_MIN_BACKEND_VERSION=}." + "cuDNN Attention is not supported due to insufficient cuDNN runtime version " + f"{backend_version=}, expected at least {CUDNN_MIN_BACKEND_VERSION=}." ) return False - if not version_at_least(cudnn.__version__, CUDNN_MIN_FRONTEND_VERSION): - log.debug( - "cuDNN Attention is not supported due to insufficient cuDNN frontend version " - f"{cudnn.__version__}, expected at least {CUDNN_MIN_FRONTEND_VERSION}." - ) + # The cuDNN SDPA ATen op is the mechanism this backend relies on; if it is missing (unexpected on + # a CUDA build), decline rather than fail later at call time. + if not hasattr(torch.ops.aten, "_scaled_dot_product_cudnn_attention"): + log.debug("cuDNN Attention is not supported because torch lacks the cuDNN SDPA ATen op.") return False return True diff --git a/cosmos_framework/model/attention/cudnn/checks.py b/cosmos_framework/model/attention/cudnn/checks.py index ffc5f96a..e6e941af 100644 --- a/cosmos_framework/model/attention/cudnn/checks.py +++ b/cosmos_framework/model/attention/cudnn/checks.py @@ -13,10 +13,83 @@ import torch from cosmos_framework.model.attention.checks import attention_param_checks, attention_tensor_checks -from cosmos_framework.model.attention.cudnn import CUDNN_DISALLOWED, CUDNN_SUPPORTED -from cosmos_framework.model.attention.cudnn.meta import get_bwd_dtypes, get_fwd_dtypes +from cosmos_framework.model.attention.cudnn import CUDNN_SUPPORTED +from cosmos_framework.model.attention.cudnn.meta import CUDNN_HEAD_DIM_ALIGNMENT, get_bwd_dtypes, get_fwd_dtypes, get_max_head_dim from cosmos_framework.model.attention.masks import CausalType -from cosmos_framework.model.attention.utils import get_arch_tag, is_torch_compiling, log_or_raise_error +from cosmos_framework.model.attention.utils import get_arch_tag, log_or_raise_error + + +def cudnn_sdpa_eligible( + query_shape: torch.Size, + key_shape: torch.Size, + value_shape: torch.Size, + arch_tag: int, + raise_error: bool = False, +) -> bool: + """ + Mirror the cuDNN SDPA eligibility constraints that the generic ``attention_tensor_checks`` does + not cover, so unsupported shapes are rejected here (and the chooser can fall back to another + backend such as NATTEN) instead of failing deep inside the raw + ``torch.ops.aten._scaled_dot_product_cudnn_attention`` operator. + + PyTorch's cuDNN eligibility logic additionally rejects: + * a key/value sequence length of 1, and + * head dims (Q/K and V) that are not a multiple of ``CUDNN_HEAD_DIM_ALIGNMENT`` (8) or that + exceed the architecture-specific maximum (see ``get_max_head_dim``, typically 128). + + This helper is intentionally tensorless and takes ``arch_tag`` directly (rather than a device) + so the eligibility rules can be unit-tested for a specific architecture (e.g. SM120) without + requiring that physical GPU. + + Parameters: + query_shape (torch.Size): Shape of 4-D query tensor (`[batch, seqlen, heads, head_dim]`). + + key_shape (torch.Size): Shape of 4-D key tensor (`[batch, seqlen_kv, heads_kv, head_dim]`). + + value_shape (torch.Size): Shape of 4-D value tensor (`[batch, seqlen_kv, heads_kv, head_dim_v]`). + + arch_tag (int): Arch tag for the target CUDA device (see ``get_arch_tag``). Example: 90 for + Hopper, 120 for workstation/consumer Blackwell. + + raise_error (bool): whether to raise an error if any checks fail, instead of just returning + False. Default is False. + + Returns: + success (bool): whether the shapes satisfy cuDNN's SDPA eligibility constraints. + """ + target_fn = partial(log_or_raise_error, raise_error=raise_error) + + # cuDNN's fused attention rejects a KV sequence length of 1. Layout is [batch, seqlen, heads, + # head_dim], so the KV sequence length is index 1 of the key shape. + seqlen_kv = key_shape[1] + if seqlen_kv == 1: + target_fn( + f"cuDNN Attention does not support a key/value sequence length of 1, got {seqlen_kv=}.", + exception=ValueError, + ) + return False + + max_head_dim = get_max_head_dim(arch_tag) + # QK and V head dims are the last dimension of each tensor. Q/K head-dim equality and (absent + # MLA) Q/V head-dim equality are already enforced by attention_tensor_checks, but we validate + # both explicitly for precise, actionable error messages. + for name, head_dim in (("QK", query_shape[-1]), ("V", value_shape[-1])): + if head_dim % CUDNN_HEAD_DIM_ALIGNMENT != 0: + target_fn( + f"cuDNN Attention requires the {name} head dim to be a multiple of " + f"{CUDNN_HEAD_DIM_ALIGNMENT}, got {head_dim=}.", + exception=ValueError, + ) + return False + if head_dim > max_head_dim: + target_fn( + f"cuDNN Attention on this architecture ({arch_tag=}) supports a maximum {name} head " + f"dim of {max_head_dim}, got {head_dim=}.", + exception=ValueError, + ) + return False + + return True def cudnn_attention_check( @@ -76,21 +149,14 @@ def cudnn_attention_check( ) return False - if CUDNN_DISALLOWED: - target_fn("cuDNN backend is disabled. Please choose another backend.", exception=RuntimeError) - return False - if deterministic: target_fn("cuDNN Attention does not support deterministic mode.", exception=RuntimeError) return False - if is_torch_compiling(): - target_fn( - "cuDNN backend does not support torch.compile yet.", - exception=RuntimeError, - ) - return False - + # cuDNN Attention supports both the forward (inference) and backward (training) passes: it runs on + # native, differentiable PyTorch ops (F.scaled_dot_product_attention / the cuDNN SDPA ATen op), so + # autograd handles the backward pass. attention_tensor_checks validates the backward dtype below + # (see get_bwd_dtypes) whenever operands require grad. arch_tag = get_arch_tag(device) fwd_dtypes = get_fwd_dtypes(arch_tag) bwd_dtypes = get_bwd_dtypes(arch_tag) @@ -103,13 +169,27 @@ def cudnn_attention_check( supported_dtypes_forward=fwd_dtypes, supported_dtypes_backward=bwd_dtypes, supports_mla=False, - supports_gqa_mqa=False, + supports_gqa_mqa=True, raise_error=raise_error, backend_name="cuDNN Attention", ): target_fn("cuDNN does not support the given inputs.", exception=RuntimeError) return False + # Mirror PyTorch's cuDNN SDPA eligibility constraints that attention_tensor_checks does not + # cover (KV sequence length 1, head-dim alignment, and the architecture-specific head-dim + # maximum). cuDNN is first in the Blackwell backend order, so without this the chooser would + # accept these shapes and then fail inside the raw cuDNN ATen operator instead of falling back + # to another backend (e.g. NATTEN). + if not cudnn_sdpa_eligible( + query_shape=query_shape, + key_shape=key_shape, + value_shape=value_shape, + arch_tag=arch_tag, + raise_error=raise_error, + ): + return False + if is_varlen: target_fn("Varlen for cuDNN Attention is not integrated yet.", exception=RuntimeError) return False diff --git a/cosmos_framework/model/attention/cudnn/functions.py b/cosmos_framework/model/attention/cudnn/functions.py index 6d06c59c..5a2d943b 100644 --- a/cosmos_framework/model/attention/cudnn/functions.py +++ b/cosmos_framework/model/attention/cudnn/functions.py @@ -7,147 +7,57 @@ cuDNN Backend: intermediate APIs Only safe to import when CUDNN_SUPPORTED is True. -""" -import time -from functools import partial +This backend runs cuDNN attention through PyTorch's *own* cuDNN SDPA dispatch rather than the +standalone cuDNN Python frontend (``import cudnn`` + ``cudnn.pygraph``). It always uses +``torch.ops.aten._scaled_dot_product_cudnn_attention`` -- the exact ATen op that +``torch.nn.functional.scaled_dot_product_attention`` lowers to for the cuDNN backend -- and simply +discards the logsumexp when ``return_lse=False``. Unlike SDPA's public API, this op also exposes the +logsumexp statistics needed for the ``return_lse=True`` path. + +It is a first-class, meta-registered PyTorch op, so it traces cleanly under +``torch.compile(fullgraph=True)``. This deliberately replaces the previous cuDNN-frontend graph +builder, which had to be hidden behind an opaque custom op (Dynamo cannot trace the frontend's +sourceless enum objects) and required a special TorchInductor flag to run on Thor -- a flag that +regressed other GPUs. +""" import torch from torch import Tensor -from torch.amp import custom_bwd, custom_fwd -from torch.autograd import Function from cosmos_framework.model.attention.checks import assert_universal_tensor_checks from cosmos_framework.model.attention.cudnn.checks import cudnn_attention_check -from cosmos_framework.model.attention.cudnn.cudnn_forward import ( - cudnn_sdpa_fwd_generate_op, - cudnn_sdpa_fwd_generate_operands, - cudnn_sdpa_fwd_post_process, -) from cosmos_framework.model.attention.masks import CausalType -from cosmos_framework.model.attention.utils.safe_ops import log - -amp_fwd = partial(custom_fwd, device_type="cuda") -amp_bwd = partial(custom_bwd, device_type="cuda") - - -CUDNN_PADDING_REQUIRED = False - - -class CudnnAttentionAutogradFn(Function): - @staticmethod - @amp_fwd - def forward( - ctx, - query: Tensor, - key: Tensor, - value: Tensor, - num_heads: int, - is_causal: bool, - scale: float, - ) -> tuple[Tensor, Tensor]: - query = query.contiguous() - key = key.contiguous() - value = value.contiguous() - - seqlen_Q = None - seqlen_KV = None - padding_Q = 0 - padding_KV = 0 - - if CUDNN_PADDING_REQUIRED: - Q_multiplier = 256 - KV_multiplier = 256 - - if query.shape[1] % Q_multiplier != 0: - seqlen_Q = query.shape[1] - padding_Q = Q_multiplier - (seqlen_Q % Q_multiplier) - - old_shape = query.shape - query = torch.nn.functional.pad(query, (0, 0, 0, 0, 0, padding_Q), "constant", 0) - log.debug(f"cuDNN Attention: padded query from {old_shape} to {query.shape}.") - - if key.shape[1] % KV_multiplier != 0: - seqlen_KV = key.shape[1] - padding_KV = KV_multiplier - (seqlen_KV % KV_multiplier) - - old_shape = key.shape - key = torch.nn.functional.pad(key, (0, 0, 0, 0, 0, padding_KV), "constant", 0) - value = torch.nn.functional.pad(value, (0, 0, 0, 0, 0, padding_KV), "constant", 0) - log.debug(f"cuDNN Attention: padded KV from {old_shape} to {key.shape}.") - - # Transform operands to cuDNN-compatible layouts, make output tensors - (q_cudnn_layout, k_cudnn_layout, v_cudnn_layout, output_cudnn_layout, lse_cudnn_layout) = ( - cudnn_sdpa_fwd_generate_operands(q=query, k=key, v=value, num_heads=num_heads, return_lse=True) - ) - - # Construct graph - assert q_cudnn_layout.device == k_cudnn_layout.device == v_cudnn_layout.device == output_cudnn_layout.device - assert q_cudnn_layout.dtype == k_cudnn_layout.dtype == v_cudnn_layout.dtype == output_cudnn_layout.dtype - cudnn_graph_gen_start = time.time() * 1e3 - cudnn_sdpa = cudnn_sdpa_fwd_generate_op( - dtype=q_cudnn_layout.dtype, - device=q_cudnn_layout.device, - q_shape=q_cudnn_layout.shape, - q_stride=q_cudnn_layout.stride(), - k_shape=k_cudnn_layout.shape, - k_stride=k_cudnn_layout.stride(), - v_shape=v_cudnn_layout.shape, - v_stride=v_cudnn_layout.stride(), - output_shape=output_cudnn_layout.shape, - output_stride=output_cudnn_layout.stride(), - lse_shape=None if lse_cudnn_layout is None else lse_cudnn_layout.shape, - lse_stride=None if lse_cudnn_layout is None else lse_cudnn_layout.stride(), - is_causal=is_causal, - attn_scale=scale, - seqlen_Q=seqlen_Q, - seqlen_KV=seqlen_KV, - ) - cudnn_graph_gen_time = time.time() * 1e3 - cudnn_graph_gen_start - log.debug(f"cuDNN Attention forward graph generation took {cudnn_graph_gen_time:.1f} ms.") - - # Execute graph - cudnn_sdpa( - q=q_cudnn_layout, - k=k_cudnn_layout, - v=v_cudnn_layout, - output=output_cudnn_layout, - lse=lse_cudnn_layout, - ) - - # Transform outputs back to torch contiguous layouts - output, logsumexp = cudnn_sdpa_fwd_post_process( - output_cudnn_layout=output_cudnn_layout, - lse_cudnn_layout=lse_cudnn_layout, - ) - - ctx.save_for_backward(q_cudnn_layout, k_cudnn_layout, v_cudnn_layout, lse_cudnn_layout, output_cudnn_layout) - ctx.num_heads = num_heads - ctx.scale = scale - - if padding_Q > 0: - old_shape = output.shape - output = output[:, :seqlen_Q, :, :] - logsumexp = logsumexp[:, :seqlen_Q, :, :] - assert output.shape[1] == seqlen_Q - assert logsumexp.shape[1] == seqlen_Q - log.debug(f"cuDNN Attention: unpadded output from {old_shape} to {output.shape}.") - - return output, logsumexp - - @staticmethod - @amp_bwd - def backward( - ctx, grad_out: Tensor, grad_lse: Tensor - ) -> tuple[ - Tensor, - Tensor, - Tensor, - None, - None, - None, - ]: - raise NotImplementedError() + + +def _cudnn_sdpa_with_lse(q: Tensor, k: Tensor, v: Tensor, is_causal: bool, scale: float) -> tuple[Tensor, Tensor]: + """cuDNN SDPA returning logsumexp, via torch's native cuDNN ATen op. + + ``torch.ops.aten._scaled_dot_product_cudnn_attention`` is the op ``F.scaled_dot_product_attention`` + dispatches to for the cuDNN backend; unlike the public API it exposes the logsumexp statistics. The + op is autograd-aware, so backward is provided automatically (no explicit autograd wrapper here). + The op requires equal Q/KV head counts, so the caller must expand K/V heads to match Q beforehand + (GQA/MQA). Inputs are heads-first ``[B, H, S, D]``; returns output ``[B, H, S, Dv]`` and logsumexp + ``[B, H, S]`` (float32). Note some torch versions return logsumexp with a trailing singleton + dim (``[B, H, S, 1]``); callers must normalize the rank. + + Only positional args ``(query, key, value, attn_bias, compute_log_sumexp, dropout_p, is_causal, + return_debug_mask)`` plus the keyword-only ``scale`` are used; these indices/names have been + stable across torch versions, and only the first two outputs (output, logsumexp) are consumed. + """ + results = torch.ops.aten._scaled_dot_product_cudnn_attention( + q, + k, + v, + None, # attn_bias + True, # compute_log_sumexp + 0.0, # dropout_p + is_causal, + False, # return_debug_mask + scale=scale, + ) + output, logsumexp = results[0], results[1] # [B,H,S,Dv], [B,H,S] or [B,H,S,1] + return output, logsumexp def cudnn_attention( @@ -167,7 +77,7 @@ def cudnn_attention( ) -> Tensor | tuple[Tensor, Tensor]: """ Runs cuDNN Attention on given operands (Q, K, V) with the heads-last contiguous layout - (`[batch, seqlen, heads, head_dim]`). + (`[batch, seqlen, heads, head_dim]`), dispatched through PyTorch's built-in cuDNN SDPA. Parameters: query (Tensor): 4-D query tensor, with the heads-last contiguous layout @@ -182,46 +92,46 @@ def cudnn_attention( is_causal (bool): whether or not causal masking is enabled. Default is False. causal_type (CausalType): causal masking mode. Choices: `CausalType.TopLeft`, - `CausalType.BottomRight`. Required when `is_causal = True`. + `CausalType.DontCare`. Required when `is_causal = True`. cuDNN SDPA's causal masking is + top-left aligned. scale (float | None): Dot product scale (attention scale). Defaults to head_dim ** -0.5. - cumulative_seqlen_Q (Tensor | None): (varlen) Optional 1-D tensor with size `batch + 1` - indicating the cumulative sum of number of query tokens in each batch, with an - additional 0 element in the beginning. Must be passed together with - `cumulative_seqlen_KV` and `max_seqlen_{Q,KV}`. + cumulative_seqlen_Q (Tensor | None): (varlen) Not supported by this backend. - cumulative_seqlen_KV (Tensor | None): (varlen) Optional 1-D tensor with size `batch + 1` - indicating the cumulative sum of number of key/value tokens in each batch, with an - additional 0 element in the beginning. Must be passed together with - `cumulative_seqlen_Q` and `max_seqlen_{Q,KV}`. + cumulative_seqlen_KV (Tensor | None): (varlen) Not supported by this backend. - max_seqlen_Q (int | None): (varlen) Optional integer indicating the maximum query - sequence length in all batches. Must be passed together with `cumulative_seqlen_{Q,KV}` - and `max_seqlen_KV`. + max_seqlen_Q (int | None): (varlen) Not supported by this backend. - max_seqlen_KV (int | None): (varlen) Optional integer indicating the maximum key/value - sequence length in all batches. Must be passed together with `cumulative_seqlen_{Q,KV}` - and `max_seqlen_Q`. + max_seqlen_KV (int | None): (varlen) Not supported by this backend. Other Parameters: return_lse (bool): Whether to return the logsumexp values. Default is False. - backend_kwargs (dict | None): Key-value pair for passing arguments specific to cuDNN's - attention operator, if any. + backend_kwargs (dict | None): Key-value pair for passing backend-specific arguments. Only + ``deterministic`` is recognized (and must be False); any other key raises an error. - deterministic (bool): Deterministic backward pass required. + deterministic (bool): Deterministic backward pass required. Not supported by this backend. Returns: output (Tensor): 4-D output tensor, with the heads-last contiguous layout (`[batch, seqlen, heads, head_dim_v]`). - logsumexp (Tensor): logsumexp tensor, with the heads-last contiguous layout - (`[batch, seqlen, heads, 1]`). Only returned when return_lse is True. + logsumexp (Tensor): logsumexp tensor, with the heads-last layout + (`[batch, seqlen, heads]`). Only returned when return_lse is True. + NOTE: not guaranteed to be contiguous (it is a transposed view) and must not be + made contiguous, so its results stay correct when merged via `merge_attentions`. """ is_varlen = cumulative_seqlen_Q is not None assert_universal_tensor_checks(query, key, value) + + backend_kwargs = backend_kwargs.copy() if backend_kwargs is not None else {} + # Determinism in backend_kwargs supersedes primary flag, if set to True + if "deterministic" in backend_kwargs: + deterministic = deterministic or backend_kwargs["deterministic"] + del backend_kwargs["deterministic"] + assert cudnn_attention_check( query_shape=query.shape, key_shape=key.shape, @@ -236,21 +146,52 @@ def cudnn_attention( raise_error=True, ) - assert not is_varlen # cudnn_attention_check should prevent this assertion failing - - num_heads = query.shape[-2] - scale = scale if scale is not None else query.shape[-1] ** -0.5 + # cudnn_attention_check should prevent this assertion failing (varlen is not integrated). + assert not is_varlen - output, lse = CudnnAttentionAutogradFn.apply( - query, - key, - value, - num_heads, - is_causal, - scale, - ) + # cuDNN SDPA takes no extra operator arguments; reject anything unrecognized instead of silently + # ignoring it. + if backend_kwargs: + raise ValueError(f"cuDNN Attention backend received unsupported backend_kwargs: {sorted(backend_kwargs)}.") - if return_lse: - return output, lse + scale = scale if scale is not None else query.shape[-1] ** -0.5 - return output + heads = query.shape[-2] + heads_kv = key.shape[-2] + assert heads % heads_kv == 0 + group_size = heads // heads_kv + + # The cuDNN ATen op expects the heads-first layout: [B,H,S,D]. + q = query.transpose(1, 2) # [B,H,S_q,D] + k = key.transpose(1, 2) # [B,H_kv,S_kv,D] + v = value.transpose(1, 2) # [B,H_kv,S_kv,D_v] + + # Always take the LSE-capable ATen path (the same op SDPA lowers to for cuDNN) and simply drop + # the logsumexp when it isn't requested. This keeps a single code path for both cases. Tradeoff: + # the raw ATen op requires equal Q/KV head counts, so GQA/MQA must materialize expanded K/V here, + # whereas the public SDPA API (enable_gqa=True) can handle grouped heads without materializing. + if group_size > 1: + k = k.repeat_interleave(group_size, dim=1, output_size=heads) # [B,H,S_kv,D] + v = v.repeat_interleave(group_size, dim=1, output_size=heads) # [B,H,S_kv,D_v] + + output, logsumexp = _cudnn_sdpa_with_lse(q, k, v, is_causal, scale) # [B,H,S_q,D_v], [B,H,S_q(,1)] + + assert output.dim() == 4 + output = output.transpose(1, 2).contiguous() # [B,S_q,H,D_v] + + if not return_lse: + return output + + # The cuDNN ATen op returns logsumexp with a trailing singleton dim on some torch + # versions ([B,H,S_q,1]); drop it so LSE is rank-3, matching every other backend and + # what `merge_attentions` requires. Guarded on rank so we neither crash on torch builds + # that already return [B,H,S_q] nor squeeze the heads dim when S_q/H is 1. + if logsumexp.dim() == 4: + logsumexp = logsumexp.squeeze(-1) # [B,H,S_q] + + # NOTE: Do NOT call .contiguous() on LSE. Attention merging's backward pass requires the + # output and LSE tensors passed into `merge_attentions` to share the same (transposed) + # data layout; forcing contiguity here makes that backward pass incorrect (see flash2). + logsumexp = logsumexp.transpose(1, 2) # [B,S_q,H] + assert logsumexp.dim() == 3 + return output, logsumexp diff --git a/cosmos_framework/model/attention/cudnn/meta.py b/cosmos_framework/model/attention/cudnn/meta.py index 5cd15344..c4e2092f 100644 --- a/cosmos_framework/model/attention/cudnn/meta.py +++ b/cosmos_framework/model/attention/cudnn/meta.py @@ -13,6 +13,38 @@ from cosmos_framework.model.attention.utils.safe_ops import log +# cuDNN's fused attention (SDPA) requires every head dim (Q/K and V) to be a multiple of 8. This +# mirrors PyTorch's own cuDNN SDPA eligibility logic (see check_cudnn_head_dim in sdp_utils.cpp). +CUDNN_HEAD_DIM_ALIGNMENT = 8 + + +def get_max_head_dim(arch_tag: int) -> int: + """ + Returns the maximum head dim cuDNN's fused attention (SDPA) supports for the given arch. + + PyTorch's cuDNN SDPA eligibility logic enforces an architecture-specific head-dim maximum + (typically 128). We stay conservative: when unsure we prefer the lower bound so the chooser + falls back to another backend (e.g. NATTEN) rather than accepting a shape that then fails + inside the raw ``torch.ops.aten._scaled_dot_product_cudnn_attention`` operator. + + Parameters: + arch_tag (int): Arch tag for the current CUDA device. Example: 80 for A100, 90 for H100, + 120 for workstation/consumer Blackwell (e.g. RTX PRO 6000). + + Returns: + max_head_dim (int): the maximum supported head dim, or 0 if the arch is unsupported. + """ + if arch_tag < 80: + log.debug("cuDNN Attention is not supported because compute capability is below the minimum (8.0).") + return 0 + + # Hopper's cuDNN fused attention supports head dims up to 256 (FP16/BF16). Every other + # currently-supported arch (Ampere/Ada and the Blackwell family, including workstation/consumer + # sm_120/121) is held to the conservative 128 maximum until larger dims are verified. + if arch_tag == 90: + return 256 + return 128 + def get_fwd_dtypes(arch_tag: int) -> list[torch.dtype]: """ @@ -30,6 +62,7 @@ def get_fwd_dtypes(arch_tag: int) -> list[torch.dtype]: log.debug("cuDNN Attention is not supported because compute capability is below the minimum (8.0).") return [] + # As of version 91400 FP8 inference via the python frontend does not seem to work. log.debug(f"cuDNN Attention only supports FP16 and BF16 for {arch_tag=}.") return [torch.float16, torch.bfloat16] @@ -46,4 +79,10 @@ def get_bwd_dtypes(arch_tag: int) -> list[torch.dtype]: """ - return [] + if arch_tag < 80: + log.debug("cuDNN Attention is not supported because compute capability is below the minimum (8.0).") + return [] + + # cuDNN SDPA backward supports the same FP16/BF16 precisions as the forward pass. + log.debug(f"cuDNN Attention backward supports FP16 and BF16 for {arch_tag=}.") + return [torch.float16, torch.bfloat16] diff --git a/cosmos_framework/model/attention/frontend.py b/cosmos_framework/model/attention/frontend.py index 018932a8..dfc93863 100644 --- a/cosmos_framework/model/attention/frontend.py +++ b/cosmos_framework/model/attention/frontend.py @@ -21,6 +21,7 @@ universal_tensor_checks, varlen_tensor_checks, ) +from cosmos_framework.model.attention.cudnn import cudnn_attention from cosmos_framework.model.attention.flash2 import flash2_attention from cosmos_framework.model.attention.flash3 import flash3_attention from cosmos_framework.model.attention.masks import CausalType @@ -28,9 +29,9 @@ from cosmos_framework.model.attention.utils.environment import filter_attention_merge_backends from cosmos_framework.model.attention.utils.safe_ops import log - # Map backend names to their frontend attention API BACKEND_MAP = { + "cudnn": cudnn_attention, "natten": natten_attention, "flash2": flash2_attention, "flash3": flash3_attention, diff --git a/cosmos_framework/model/attention/utils/__init__.py b/cosmos_framework/model/attention/utils/__init__.py index 67eaaae5..e7aece57 100644 --- a/cosmos_framework/model/attention/utils/__init__.py +++ b/cosmos_framework/model/attention/utils/__init__.py @@ -60,6 +60,14 @@ def is_blackwell_dc(device: torch.device | None = None) -> bool: return get_arch_tag(device) in [100, 103] +def is_blackwell(device: torch.device | None = None) -> bool: + """True for any Blackwell GPU: data-center (sm_100/103), Thor (sm_110, DRIVE/Jetson Thor), + or workstation/consumer (sm_120/121, e.g. RTX PRO 6000 / RTX 50-series). This matches the + non-data-center grouping ``[110, 120, 121]`` in ``get_backend_list``. Use ``is_blackwell_dc`` + when a check must be restricted to the data-center parts (e.g. kernels only qualified there).""" + return get_arch_tag(device) in [100, 103, 110, 120, 121] + + __all__ = [ "get_arch_tag", "log_or_raise_error", @@ -68,6 +76,7 @@ def is_blackwell_dc(device: torch.device | None = None) -> bool: "is_fp8", "is_hopper", "is_blackwell_dc", + "is_blackwell", "is_torch_compiling", "torch_deterministic_mode", ] diff --git a/cosmos_framework/model/generator/mot/attention.py b/cosmos_framework/model/generator/mot/attention.py index f0d87660..7eb7fa6c 100644 --- a/cosmos_framework/model/generator/mot/attention.py +++ b/cosmos_framework/model/generator/mot/attention.py @@ -111,34 +111,74 @@ def two_way_attention( causal_v, _ = get_causal_seq(packed_value_states) full_q, full_q_offsets = get_full_only_seq(packed_query_states) + # NOTE: we can only use the don't care causal mask when we know seqlen_Q == seqlen_KV. + # Since this is a varlen use case, we would need to statically check all Q and KV offsets + # are the same. + # We don't want to launch a kernel just to perform this check and slow down our model, and + # we definitely don't want to complicate the sequence_packing code so that it performs a + # static check when creating the packed sequence and metadata. Instead, we just rely + # on causal_q_offsets and causal_k_offsets being the same tensor. + use_dont_care_mask = causal_q_offsets is causal_k_offsets + sample_offsets = packed_query_states["sample_offsets"] - use_dont_care_mask = causal_q_offsets is causal_k_offsets + # Number of packed samples. sample_offsets has shape [num_samples + 1], so this is a static + # (metadata) shape read: it does not launch a kernel or force a device sync. + num_samples = sample_offsets.shape[0] - 1 + + # With a single sample there is exactly one sequence in the pack, so the varlen (sequence-packed) + # metadata is redundant and we can call the dense attention API instead (no cumulative/max seqlen + # args). This remains correct in the presence of trailing padding: for causal self-attention the + # mask never lets a real query attend to padded keys (padding is appended after all real tokens), + # get_all_seq returns unpadded KV for the full path, and any padded query rows are independent of + # the real rows and simply discarded downstream. + # + # The dense path is gated to forward-only (inference) execution via + # torch.is_grad_enabled(), which is False under torch.no_grad()/torch.inference_mode() + # and True during training. This avoids branching on the sample count during training, + # where batch composition varies between single- and multi-sample packs; keeping a single + # code path there prevents torch.compile from specializing on both shapes and incurring + # the associated recompilation overhead. + use_dense = num_samples == 1 and not torch.is_grad_enabled() + + if use_dense: + causal_varlen_kwargs = {} + else: + causal_varlen_kwargs = dict( + cumulative_seqlen_Q=causal_q_offsets, + cumulative_seqlen_KV=causal_k_offsets, + max_seqlen_Q=packed_query_states["max_causal_len"], + max_seqlen_KV=packed_query_states["max_causal_len"], + ) # NOTE: cosmos_framework attention is BSHD in, BSHD out causal_res = attention( causal_q.unsqueeze(0), # [1,N_und,heads,head_dim] causal_k.unsqueeze(0), # [1,N_und,heads,head_dim] causal_v.unsqueeze(0), # [1,N_und,heads,head_dim] - cumulative_seqlen_Q=causal_q_offsets, - cumulative_seqlen_KV=causal_k_offsets, - max_seqlen_Q=packed_query_states["max_causal_len"], - max_seqlen_KV=packed_query_states["max_causal_len"], is_causal=True, causal_type=CausalType.DontCare if use_dont_care_mask else CausalType.TopLeft, + **causal_varlen_kwargs, ) # [1,N_und,heads,head_dim] # [1,N_und,heads,head_dim] -> [N_und,heads,head_dim] -> [N_und,heads*head_dim] causal_out = causal_res.squeeze(0).flatten(-2, -1) # type: ignore # [N_und,heads*head_dim] + if use_dense: + full_varlen_kwargs = {} + else: + full_varlen_kwargs = dict( + cumulative_seqlen_Q=full_q_offsets, + cumulative_seqlen_KV=sample_offsets, + max_seqlen_Q=packed_query_states["max_full_len"], + max_seqlen_KV=packed_query_states["max_sample_len"], + ) + full_res = attention( full_q.unsqueeze(0), # [1,N_full,heads,head_dim] get_all_seq(packed_key_normalized).unsqueeze(0), # [1,N_all,heads,head_dim] normed und K for gen get_all_seq(packed_value_states).unsqueeze(0), # [1,N_all,heads,head_dim] - cumulative_seqlen_Q=full_q_offsets, - cumulative_seqlen_KV=sample_offsets, - max_seqlen_Q=packed_query_states["max_full_len"], - max_seqlen_KV=packed_query_states["max_sample_len"], + **full_varlen_kwargs, ) # [1,N_full,heads,head_dim] # [1,N_full,heads,head_dim] -> [N_full,heads,head_dim] -> [N_full,heads*head_dim] diff --git a/cosmos_framework/model/generator/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py index 656364b2..6fed5dee 100644 --- a/cosmos_framework/model/generator/omni_mot_model.py +++ b/cosmos_framework/model/generator/omni_mot_model.py @@ -392,13 +392,9 @@ def set_up_memory(self) -> None: def set_up_parallelism(self) -> None: """Set up the fsdp for the model.""" - if not torch.distributed.is_initialized(): - self.parallel_dims = None - return - self.parallel_dims = ParallelDims( enable_inference_mode=self.config.parallelism.enable_inference_mode, - world_size=torch.distributed.get_world_size(), + world_size=torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1, dp_shard=self.config.parallelism.data_parallel_shard_degree, cfgp=self.config.parallelism.cfg_parallel_shard_degree, cp=self.config.parallelism.context_parallel_shard_degree, @@ -1666,8 +1662,13 @@ def _prepare_inference_data( ) # 2. Get data and condition (same as training) - # This encodes vision to x0_tokens - gen_data_clean = self.get_data_and_condition(data_batch) + # This encodes vision to x0_tokens. Pass each sample's vision conditioning + # frame indexes so a causal tokenizer can skip encoding pixel frames that only + # feed generated (non-conditioned) latent positions (e.g. policy / + # forward-dynamics condition latent frame 0 only, so only the first pixel frame + # is encoded instead of the whole clip). + vision_condition_indexes = [plan.condition_frame_indexes_vision for plan in sequence_plans] + gen_data_clean = self.get_data_and_condition(data_batch, vision_condition_indexes=vision_condition_indexes) num_items_per_sample = gen_data_clean.num_vision_items_per_sample # None for standard T2I/T2V @@ -2846,11 +2847,139 @@ def validation_step(self, data_batch: dict[str, torch.Tensor], iteration: int): def forward(self, xt, t): pass - def get_data_and_condition(self, data_batch: dict[str, torch.Tensor], iteration: int = 1) -> GenerationDataClean: + def _encode_vision_x0_tokens( + self, + raw_state_vision: list[torch.Tensor], + num_vision_items_per_sample: list[int] | None, + vision_condition_indexes: list[list[int]] | None, + ) -> list[torch.Tensor]: + """Encode vision items into x0 latent tokens, optionally skipping unused frames. + + Default behavior (``vision_condition_indexes is None``) encodes every pixel + frame of every vision item. This is the path used during training and by any + caller that does not opt in, and it is byte-for-byte the original encode loop. + + Inference optimization: when ``vision_condition_indexes`` is provided (one + conditioning-frame-index list per sample, from each ``SequencePlan``) and the + vision tokenizer is temporally causal, only the pixel-frame prefix required to + reconstruct the highest conditioned latent frame is encoded. The remaining + latent frames are generated (pure-noise) positions that the ``condition_mask`` + blend discards, so they are zero-filled instead of encoded. Under a causal VAE + the kept latents are identical to the corresponding frames of a full-clip + encode, while the encode processes far fewer frames (e.g. 1 pixel frame instead + of the full clip for policy / forward-dynamics modes that only condition latent + frame 0). Inverse-dynamics conditions every latent frame, so it keeps the full + encode automatically. + + The optimization falls back to a full encode for multi-vision samples, + non-causal tokenizers, samples with no conditioning frames, and any item whose + full clip is already the minimal prefix. + """ + + def _full_encode(state: torch.Tensor) -> torch.Tensor: + return self.encode(state).contiguous().float() + + # Only opt in when a caller supplied per-sample conditioning indexes, the + # samples map 1:1 to vision items (no multi-vision flattening), and the + # tokenizer is causal (so a pixel prefix reproduces the leading latents). + optimization_applicable = ( + vision_condition_indexes is not None + and num_vision_items_per_sample is None + and self.tokenizer_vision_gen is not None + and self.tokenizer_vision_gen.is_causal + and len(vision_condition_indexes) == len(raw_state_vision) + ) + if not optimization_applicable: + return [_full_encode(raw_state_vision_i) for raw_state_vision_i in raw_state_vision] + + # The prefix-encode optimization zero-fills the generated (pure-noise) latent + # frames that the condition_mask blend later discards. That is only valid when + # no gradient is required: torch.is_grad_enabled() is False under both + # torch.inference_mode() and torch.no_grad(), so this asserts we are on an + # inference path and never silently corrupts a training forward. + if torch.is_grad_enabled(): + raise ValueError( + "prefix-encode optimization is inference-only, but grad is enabled " + "(expected torch.inference_mode()/torch.no_grad())" + ) + + tokenizer = self.tokenizer_vision_gen + assert vision_condition_indexes is not None # narrowed by optimization_applicable above + x0_tokens_vision: list[torch.Tensor] = [] + for raw_state_vision_i, condition_indexes in zip(raw_state_vision, vision_condition_indexes, strict=True): + # The temporal axis is the first of the trailing (T, H, W) dims for both the + # [B, C, T, H, W] and [C, T, H, W] layouts the tokenizer accepts, and encode + # preserves rank so the same index locates T in the latent output. + temporal_dim = raw_state_vision_i.ndim - 3 + num_pixel_frames = int(raw_state_vision_i.shape[temporal_dim]) + num_latent_frames = tokenizer.get_latent_num_frames(num_pixel_frames) + + valid_condition = [idx for idx in condition_indexes if 0 <= idx < num_latent_frames] + if not valid_condition: + # No conditioning frames to preserve (e.g. fully generated T2V item): + # keep the full encode so behavior is unchanged for non-action items. + x0_tokens_vision.append(_full_encode(raw_state_vision_i)) + continue + + needed_latent_frames = max(valid_condition) + 1 + needed_pixel_frames = tokenizer.get_pixel_num_frames(needed_latent_frames) + assert needed_pixel_frames <= num_pixel_frames, ( + f"needed_pixel_frames ({needed_pixel_frames}) cannot be greater than " + f"num_pixel_frames ({num_pixel_frames})" + ) + if needed_pixel_frames == num_pixel_frames: + # The conditioning already spans (nearly) the whole clip; no work saved. + x0_tokens_vision.append(_full_encode(raw_state_vision_i)) + continue + + prefix = raw_state_vision_i.narrow(temporal_dim, 0, needed_pixel_frames) + prefix_latent = _full_encode(prefix) # leading latent frames, causally exact + + # The scatter below assumes the pixel->latent round trip is exact, i.e. that + # encoding ``needed_pixel_frames`` yields exactly ``needed_latent_frames`` + # latents. All current causal tokenizers satisfy this, but guard it: an + # overshoot would make the narrow below request more frames than full_latent + # holds (confusing runtime error), and an undershoot would silently zero-fill + # the highest conditioning frame and corrupt the condition_mask blend with no + # error at all. + num_prefix_latent_frames = prefix_latent.shape[temporal_dim] + assert num_prefix_latent_frames == needed_latent_frames, ( + f"causal tokenizer pixel<->latent round trip is not exact: encoding " + f"{needed_pixel_frames} pixel frames (from get_pixel_num_frames({needed_latent_frames})) " + f"produced {num_prefix_latent_frames} latent frames, expected {needed_latent_frames}" + ) + + # Scatter the encoded prefix into a full-length latent; the unconditioned + # tail frames stay zero since the condition_mask blend replaces them with noise. + full_shape = list(prefix_latent.shape) + full_shape[temporal_dim] = num_latent_frames + full_latent = prefix_latent.new_zeros(full_shape) + full_latent.narrow(temporal_dim, 0, prefix_latent.shape[temporal_dim]).copy_(prefix_latent) + x0_tokens_vision.append(full_latent) + + return x0_tokens_vision + + def get_data_and_condition( + self, + data_batch: dict[str, torch.Tensor], + iteration: int = 1, + vision_condition_indexes: list[list[int]] | None = None, + ) -> GenerationDataClean: """ - Get raw data of different modalities from databatch - Tokenize into corresponding latents - Load other conditioning information if any (fps, etc.) + + Args: + data_batch: Raw data batch from the dataloader. + iteration: Current iteration (only used for time-based logging during training). + vision_condition_indexes: Optional per-sample list of conditioning latent + frame indexes (one list per sample, taken from each ``SequencePlan``). + When provided (inference only), it enables the causal-VAE prefix-encode + optimization in ``_encode_vision_x0_tokens`` that skips encoding pixel + frames which only feed generated (non-conditioned) latent positions. + Leaving it ``None`` (the default, used during training) encodes every + frame — the original behavior. """ # Detect whether any sample has multiple vision items (e.g. image editing). # If so, track the count per sample before all vision items from this batch are flattened into a list. @@ -2898,13 +3027,16 @@ def get_data_and_condition(self, data_batch: dict[str, torch.Tensor], iteration: if log_enc_time: timer = Timer(unit="s") timer.start() + # Vision (image/video) raw state and tokenized latent state self._normalize_video_databatch_inplace(data_batch) self._augment_image_dim_inplace(data_batch) # converts each image tensor to (1, C, 1, H, W) raw_state_vision = data_batch[self.input_image_key if is_image_batch else self.input_video_key] - x0_tokens_vision = [ - self.encode(raw_state_vision_i).contiguous().float() for raw_state_vision_i in raw_state_vision - ] + x0_tokens_vision = self._encode_vision_x0_tokens( + raw_state_vision, + num_vision_items_per_sample, + vision_condition_indexes, + ) frame_size = data_batch.get("image_size", None) if frame_size is not None: diff --git a/cosmos_framework/model/generator/tokenizers/uniae/noncausal_4x16x16.py b/cosmos_framework/model/generator/tokenizers/uniae/noncausal_4x16x16.py index a9e0dac7..c65788bb 100644 --- a/cosmos_framework/model/generator/tokenizers/uniae/noncausal_4x16x16.py +++ b/cosmos_framework/model/generator/tokenizers/uniae/noncausal_4x16x16.py @@ -17,7 +17,10 @@ recon = vae.decode(latents) # [B, 48, T_p, H//16, W//16] -> [B, 3, 4*T_p, H, W] """ +import os from collections.abc import Mapping, Sequence +from numbers import Real +from typing import Any import torch @@ -32,52 +35,174 @@ normalize_resolution_int_mapping, ) from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution +from cosmos_framework.model.tokenizer.checkpoint_io import ( + DCP_MODEL_LOAD_INFO_KEY, + DCPModelLoadInfo, + load_dcp_model_checkpoint, + load_torch_checkpoint, +) +from cosmos_framework.model.tokenizer.models.architecture import get_siglip2_so400m_common_arch from cosmos_framework.model.tokenizer.models.dense_runtime import DenseAutoencoderRuntime from cosmos_framework.model.tokenizer.models.sparse_autoencoder import AutoencoderKL -# S3 architecture config (avoids importing configs/base which pulls in loss deps) _S3_ARCH = dict( - patch_size=(4, 16, 16), - in_channels=3072, - out_channels=3072, - # Encoder - encoder_model_channels=1152, - encoder_num_blocks=27, - encoder_num_heads=16, - encoder_mlp_channels=4304, - encoder_pe_mode="joint", - encoder_qk_rms_norm=False, - encoder_use_bias=True, - encoder_use_rms_norm=False, - # Decoder - decoder_model_channels=1152, - decoder_num_blocks=27, - decoder_num_heads=16, - decoder_mlp_channels=4304, - decoder_pe_mode="joint", - decoder_qk_rms_norm=True, - decoder_use_bias=False, - decoder_use_rms_norm=True, - # Common settings - use_decoder=True, - quantizer_type="rq", - quantizer_codebook_size=65536, - quantizer_num_codebooks=1, - quantizer_chunk_size=1, - use_vf_loss=False, - freeze_encoder=False, - pretrained_model_name="google/siglip2-so400m-patch16-naflex", - concat_latent=None, - random_num_sample_frames_batch_sizes=[8, 12, 16, 20, 24], - inference_num_sample_frames_batch_size=16, - inference_num_sample_frames_stride=16, - inference_kv_cache_size=0, + **get_siglip2_so400m_common_arch(), use_quantizer=False, use_dual_latent=False, use_text_alignment=False, use_post_text_alignment=False, ) +_IGNORED_LEGACY_CHECKPOINT_PREFIXES = ( + "ema.", + "loss_fn.", + "text_decoder_wrapper.", + "tokenizer.", +) +_LEGACY_CHECKPOINT_SUFFIXES = (".pt", ".pth", ".ckpt") + + +def _is_legacy_checkpoint_file(checkpoint_path: str) -> bool: + """Distinguish torch checkpoint files from DCP component directories.""" + is_remote = checkpoint_path.startswith("s3://") + if not is_remote and os.path.isdir(checkpoint_path): + return False + return checkpoint_path.lower().endswith(_LEGACY_CHECKPOINT_SUFFIXES) or ( + not is_remote and os.path.isfile(checkpoint_path) + ) + + +def _derive_legacy_latent_norm_path(checkpoint_path: str) -> str: + """Derive the latent-stat sidecar for a legacy checkpoint file.""" + checkpoint_stem, checkpoint_suffix = os.path.splitext(checkpoint_path) + if checkpoint_suffix.lower() in _LEGACY_CHECKPOINT_SUFFIXES: + return checkpoint_stem + "_latent_norm.pt" + return checkpoint_path + "_latent_norm.pt" + + +def _coerce_latent_norm_vector( + values: object, + *, + field_name: str, + z_dim: int, + source: str, + strictly_positive: bool = False, +) -> torch.Tensor: + """Convert one tensor/JSON vector to validated CPU float64 statistics.""" + if isinstance(values, torch.Tensor): + if values.dtype == torch.bool: + raise ValueError(f"Latent-normalization {field_name} in {source} must contain real numbers, not booleans.") + if values.is_complex(): + raise ValueError( + f"Latent-normalization {field_name} in {source} must contain real numbers, not complex values." + ) + vector = values.detach().to(device="cpu", dtype=torch.float64) # [C_candidate] + elif isinstance(values, Sequence) and not isinstance(values, (str, bytes)): + if any(isinstance(value, bool) or not isinstance(value, Real) for value in values): + raise ValueError(f"Latent-normalization {field_name} in {source} must contain only real numbers.") + vector = torch.tensor(list(values), dtype=torch.float64) # [C_candidate] + else: + raise TypeError( + f"Latent-normalization {field_name} in {source} must be a tensor or numeric sequence, " + f"got {type(values).__name__}." + ) + + if vector.ndim != 1 or vector.shape[0] != z_dim: + raise ValueError( + f"Latent-normalization {field_name} in {source} must have shape ({z_dim},), got {tuple(vector.shape)}." + ) + finite_mask = torch.isfinite(vector) # [C] + if not bool(finite_mask.all().item()): + raise ValueError(f"Latent-normalization {field_name} in {source} contains non-finite values.") + if strictly_positive: + positive_mask = vector > 0 # [C] + if not bool(positive_mask.all().item()): + raise ValueError(f"Latent-normalization {field_name} in {source} must be strictly positive.") + return vector + + +def _load_latent_norm_stats( + norm_path: str, + *, + backend_args: dict[str, str] | None, + z_dim: int, + dtype: torch.dtype, + device: str, +) -> tuple[torch.Tensor, torch.Tensor]: + """Load latent statistics and return validated runtime mean and inverse standard deviation.""" + if norm_path.lower().endswith(".json"): + norm_stats = easy_io.load(norm_path, backend_args=backend_args) + else: + norm_stats = easy_io.load(norm_path, backend_args=backend_args, map_location="cpu", weights_only=True) + if not isinstance(norm_stats, Mapping): + raise TypeError( + f"Latent-normalization sidecar {norm_path} must contain a mapping, got {type(norm_stats).__name__}." + ) + + stats_z_dim = norm_stats.get("z_dim") + if stats_z_dim is not None and (isinstance(stats_z_dim, bool) or not isinstance(stats_z_dim, int)): + raise ValueError(f"Latent-normalization z_dim in {norm_path} must be an integer, got {stats_z_dim!r}.") + if isinstance(stats_z_dim, int) and stats_z_dim != z_dim: + raise ValueError(f"Latent-normalization z_dim in {norm_path} is {stats_z_dim}, expected {z_dim}.") + if "mean" not in norm_stats or "std" not in norm_stats: + raise ValueError(f"Latent-normalization sidecar {norm_path} must contain mean and std entries.") + + mean_cpu = _coerce_latent_norm_vector( + norm_stats["mean"], + field_name="mean", + z_dim=z_dim, + source=norm_path, + ) # [C] + std_cpu = _coerce_latent_norm_vector( + norm_stats["std"], + field_name="std", + z_dim=z_dim, + source=norm_path, + strictly_positive=True, + ) # [C] + mean = mean_cpu.to(dtype=dtype, device=device) # [C] + std = std_cpu.to(dtype=dtype, device=device) # [C] + mean_finite_mask = torch.isfinite(mean) # [C] + std_finite_positive_mask = torch.isfinite(std) & (std > 0) # [C] + if not bool(mean_finite_mask.all().item()): + raise ValueError(f"Latent-normalization mean in {norm_path} is non-finite after conversion to {dtype}.") + if not bool(std_finite_positive_mask.all().item()): + raise ValueError( + f"Latent-normalization std in {norm_path} must remain finite and positive after conversion to {dtype}." + ) + inv_std = std.reciprocal() # [C] + inv_std_finite_mask = torch.isfinite(inv_std) # [C] + if not bool(inv_std_finite_mask.all().item()): + raise ValueError(f"Latent-normalization inverse std in {norm_path} is non-finite after conversion to {dtype}.") + return mean, inv_std + + +def _extract_visual_tokenizer_state_dict(model_state: Mapping[str, Any]) -> dict[str, Any]: + """Extract bare AutoencoderKL state from a full TokenizerModel checkpoint.""" + network_state = { + key.removeprefix("network."): value for key, value in model_state.items() if key.startswith("network.") + } + return network_state if network_state else dict(model_state) + + +def _get_dcp_unexpected_visual_keys(load_info: DCPModelLoadInfo) -> list[str]: + """Return unexpected keys owned by the tokenizer network from a DCP report.""" + return [key.removeprefix("network.") for key in load_info.unexpected_checkpoint_keys if key.startswith("network.")] + + +def _validate_uniae_checkpoint_keys(missing: list[str], unexpected: list[str], checkpoint_path: str) -> None: + """Reject promoted checkpoints that do not fully cover the visual tokenizer.""" + if missing: + raise RuntimeError( + f"UniAE checkpoint {checkpoint_path} is missing {len(missing)} visual model keys (sample: {missing[:5]})." + ) + invalid_unexpected = [key for key in unexpected if not key.startswith(_IGNORED_LEGACY_CHECKPOINT_PREFIXES)] + if invalid_unexpected: + raise RuntimeError( + f"UniAE checkpoint {checkpoint_path} has {len(invalid_unexpected)} unexpected visual model keys " + f"(sample: {invalid_unexpected[:5]})." + ) + class UniAEVAE: """UniAE S3 VAE wrapper for diffusion training. @@ -95,6 +220,7 @@ def __init__( z_dim: int = 48, vae_pth: str = "", object_store_credential_path_pretrained: str = "", + latent_norm_path: str = "", dtype: torch.dtype = torch.bfloat16, device: str = "cuda", backend: str = "batched", @@ -102,7 +228,9 @@ def __init__( pixel_trim: bool = True, chunk_size: int | Mapping[str, int] = 16, encode_chunk_batch_size: int | Mapping[str, int] = 1, - ): + ) -> None: + if torch.device(device).type == "meta": + raise ValueError("UniAEVAE requires a concrete CPU or CUDA device; device='meta' is not supported.") chunk_size = normalize_resolution_int_mapping(chunk_size, name="chunk_size") if any(chunk_frames % 4 != 0 for chunk_frames in chunk_size.values()): raise ValueError("chunk_size values must be multiples of 4.") @@ -135,11 +263,14 @@ def __init__( if not vae_pth: raise ValueError("vae_pth must be provided to load latent normalization stats") vae_pth_str = str(vae_pth) - # Derive stats path: strip .pt suffix, append _latent_norm.pt - if vae_pth_str.endswith(".pt"): - norm_pth = vae_pth_str[:-3] + "_latent_norm.pt" + is_legacy_checkpoint = _is_legacy_checkpoint_file(vae_pth_str) + # Legacy files use a paired _latent_norm.pt sidecar. + if latent_norm_path: + norm_pth = latent_norm_path + elif is_legacy_checkpoint: + norm_pth = _derive_legacy_latent_norm_path(vae_pth_str) else: - norm_pth = vae_pth_str + "_latent_norm.pt" + raise ValueError("latent_norm_path is required when vae_pth points to a DCP checkpoint directory.") if norm_pth.startswith("s3://"): norm_backend_args = { "backend": "s3", @@ -147,13 +278,17 @@ def __init__( } else: norm_backend_args = None - norm_stats = easy_io.load(norm_pth, backend_args=norm_backend_args, map_location="cpu", weights_only=False) - mean = norm_stats["mean"].to(dtype=dtype, device=device) - std = norm_stats["std"].to(dtype=dtype, device=device) - self._latent_mean = mean.view(1, z_dim, 1, 1, 1) - self._latent_inv_std = (1.0 / std).view(1, z_dim, 1, 1, 1) - - # make compatible with meta device + mean, inv_std = _load_latent_norm_stats( + norm_pth, + backend_args=norm_backend_args, + z_dim=z_dim, + dtype=dtype, + device=device, + ) # [C],[C] + self._latent_mean = mean.view(1, z_dim, 1, 1, 1) # [1,C,1,1,1] + self._latent_inv_std = inv_std.view(1, z_dim, 1, 1, 1) # [1,C,1,1,1] + + # Construct the autoencoder on the requested concrete device. autoencoder = AutoencoderKL( **_S3_ARCH, latent_channels=z_dim, @@ -164,27 +299,44 @@ def __init__( # Load checkpoint if vae_pth and get_rank() == 0: - if str(vae_pth).startswith("s3://"): - backend_args = {"backend": "s3", "s3_credential_path": object_store_credential_path_pretrained} + dcp_unexpected_keys: list[str] = [] + if is_legacy_checkpoint: + if str(vae_pth).startswith("s3://"): + backend_args = { + "backend": "s3", + "s3_credential_path": object_store_credential_path_pretrained, + } + else: + backend_args = None + state_dict = load_torch_checkpoint( + vae_pth, + backend_args=backend_args, + map_location="cpu", + ) + if "model" in state_dict: + model_state = state_dict["model"] + elif "state_dict" in state_dict: + model_state = state_dict["state_dict"] + else: + model_state = state_dict + model_state = _extract_visual_tokenizer_state_dict(model_state) else: - backend_args = None - state_dict = easy_io.load(vae_pth, backend_args=backend_args, map_location="cpu", weights_only=False) - if "model" in state_dict: + state_dict = load_dcp_model_checkpoint( + autoencoder, + vae_pth, + checkpoint_key_prefix="network.", + s3_credential=object_store_credential_path_pretrained, + ) model_state = state_dict["model"] - elif "state_dict" in state_dict: - model_state = state_dict["state_dict"] - else: - model_state = state_dict - # Checkpoint may be saved from a wrapper with a 'network.' prefix — strip it. - if any(k.startswith("network.") for k in model_state): - model_state = { - k[len("network.") :] if k.startswith("network.") else k: v for k, v in model_state.items() - } + dcp_load_info = state_dict[DCP_MODEL_LOAD_INFO_KEY] + if not isinstance(dcp_load_info, DCPModelLoadInfo): + raise TypeError(f"Invalid DCP model load info: {type(dcp_load_info).__name__}") + dcp_unexpected_keys = _get_dcp_unexpected_visual_keys(dcp_load_info) missing, unexpected = autoencoder.load_state_dict(model_state, strict=False) - if missing: - log.warning(f"Missing keys: {len(missing)} (e.g., {missing[:3]})") - if unexpected: - log.warning(f"Unexpected keys: {len(unexpected)} (e.g., {unexpected[:3]})") + all_unexpected = [*unexpected, *dcp_unexpected_keys] + _validate_uniae_checkpoint_keys(missing, all_unexpected, str(vae_pth)) + if all_unexpected: + log.info(f"Ignored {len(all_unexpected)} non-visual checkpoint keys (e.g., {all_unexpected[:3]}).") log.info(f"Loaded checkpoint from {vae_pth}") elif vae_pth: autoencoder.to_empty(device=device) @@ -364,16 +516,25 @@ def __init__( bucket_name: str = "", object_store_credential_path_pretrained: str = "", vae_path: str = "", + latent_norm_path: str = "", encode_chunk_frames: int | Mapping[str, int] = 16, encode_chunk_batch_size: int | Mapping[str, int] = 1, spatial_compression_factor: int = 16, temporal_compression_factor: int = 4, - pad_frames: int = 0, + pad_frames: int = 1, pixel_trim: bool = True, backend: str = "batched_with_padding", causal: bool = False, - ): - super().__init__(object_store_credential_path_pretrained) + ) -> None: + if spatial_compression_factor != 16: + raise ValueError( + f"UniAEVAEInterface requires spatial_compression_factor=16, got {spatial_compression_factor}." + ) + if temporal_compression_factor != 4: + raise ValueError( + f"UniAEVAEInterface requires temporal_compression_factor=4, got {temporal_compression_factor}." + ) + super().__init__(object_store_credential_path_pretrained or None) self._causal = causal assert not self._causal, "UniAEVAEInterface is a non-causal tokenizer; causal must be False." self._spatial_compression_factor = spatial_compression_factor @@ -400,12 +561,16 @@ def __init__( self.use_streaming_encode = False vae_full_path = vae_path + latent_norm_full_path = latent_norm_path if bucket_name and not vae_path.startswith("s3://"): vae_full_path = f"s3://{bucket_name}/{vae_path}" + if bucket_name and latent_norm_path and not latent_norm_path.startswith("s3://"): + latent_norm_full_path = f"s3://{bucket_name}/{latent_norm_path}" self.vae = UniAEVAE( vae_pth=vae_full_path, object_store_credential_path_pretrained=object_store_credential_path_pretrained, + latent_norm_path=latent_norm_full_path, pad_frames=pad_frames, pixel_trim=pixel_trim, backend=backend, @@ -414,7 +579,7 @@ def __init__( ) self.is_compiled = False - def reset_dtype(self): + def reset_dtype(self) -> None: pass def encode(self, state: torch.Tensor) -> torch.Tensor: @@ -483,11 +648,11 @@ def get_latent_temporal_positions( ) @property - def spatial_compression_factor(self): + def spatial_compression_factor(self) -> int: return self._spatial_compression_factor @property - def temporal_compression_factor(self): + def temporal_compression_factor(self) -> int: return self._temporal_compression_factor @property diff --git a/cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py b/cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py index fea6320d..0117a13c 100644 --- a/cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py +++ b/cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py @@ -1285,7 +1285,10 @@ def _collect_warmup_shapes( raise ValueError(f"Aspect ratio {aspect_ratio} not found in resolution {res_key}") res_dict = {aspect_ratio: res_dict[aspect_ratio]} - for H, W in res_dict.values(): + # ``VIDEO_RES_SIZE_INFO`` stores each bucket as ``(W, H)`` (e.g. the + # landscape ``"4,3": (736, 544)``), so unpack width-first to match the + # runtime encode key, which is derived from the actual ``(H, W)`` tensor. + for W, H in res_dict.values(): if isinstance(tokenizer.encode_chunk_frames, Mapping): chunk_frames = tokenizer.encode_chunk_frames[res_key] else: diff --git a/cosmos_framework/model/tokenizer/checkpoint_io.py b/cosmos_framework/model/tokenizer/checkpoint_io.py new file mode 100644 index 00000000..1f2079a7 --- /dev/null +++ b/cosmos_framework/model/tokenizer/checkpoint_io.py @@ -0,0 +1,318 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Safe checkpoint loading helpers for tokenizer code.""" + +from __future__ import annotations + +import io +import os +import re +from collections import defaultdict +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from omegaconf.base import ContainerMetadata, Metadata +from omegaconf.dictconfig import DictConfig +from omegaconf.listconfig import ListConfig +from omegaconf.nodes import AnyNode + +from cosmos_framework.utils.easy_io import easy_io + +TorchMapLocation = str | torch.device | Callable[[Any, str], Any] | dict[str, str] | None +DCP_MODEL_LOAD_INFO_KEY = "dcp_model_load_info" +_RQ_CODEBOOK_STAGE_KEY_PATTERN = re.compile(r"(?:^|\.)quantizer\.codebooks\.(\d+)(?:\.|$)") + + +@dataclass(frozen=True) +class DCPModelLoadInfo: + """Model-key coverage discovered from one DCP component's metadata.""" + + loaded_model_keys: frozenset[str] + missing_model_keys: frozenset[str] + unexpected_checkpoint_keys: frozenset[str] + + +def infer_rq_codebook_depth(state_keys: Iterable[str], *, source: str) -> int | None: + """Infer a contiguous residual-quantizer depth from state-dict keys.""" + stage_indices = { + int(match.group(1)) + for key in state_keys + if (match := _RQ_CODEBOOK_STAGE_KEY_PATTERN.search(str(key))) is not None + } + if not stage_indices: + return None + + inferred_depth = max(stage_indices) + 1 + expected_indices = set(range(inferred_depth)) + if stage_indices != expected_indices: + raise ValueError( + f"Residual-quantizer stages in {source} are not contiguous from zero: " + f"found {sorted(stage_indices)}, expected {sorted(expected_indices)}." + ) + return inferred_depth + + +def get_model_state_keys(model: Any) -> set[str]: + """Return model-state names without constructing a mapping for torch modules.""" + if isinstance(model, nn.Module): + parameter_keys = {name for name, _param in model.named_parameters(remove_duplicate=False)} + buffer_keys = {name for name, _buffer in model.named_buffers(remove_duplicate=False)} + return parameter_keys | buffer_keys + + state_dict = getattr(model, "state_dict", None) + return set(state_dict()) if callable(state_dict) else set() + + +def validate_rq_checkpoint_depth( + model: nn.Module, + checkpoint_keys: Iterable[str], + *, + checkpoint_path: str, + model_keys: Iterable[str] | None = None, +) -> None: + """Reject RQ checkpoints whose persisted depth disagrees with the instantiated model.""" + resolved_model_keys = get_model_state_keys(model) if model_keys is None else model_keys + model_depth = infer_rq_codebook_depth(resolved_model_keys, source="the instantiated model") + checkpoint_depth = infer_rq_codebook_depth( + checkpoint_keys, + source=f"checkpoint {checkpoint_path}", + ) + if model_depth is None or checkpoint_depth is None or model_depth == checkpoint_depth: + return + + raise ValueError( + f"Residual-quantizer depth mismatch for checkpoint {checkpoint_path}: " + f"the instantiated model has {model_depth} stage(s), but the checkpoint contains {checkpoint_depth}. " + "Historical UniAE configs may record quantizer_num_codebooks=1 even though their checkpoints contain " + "four stages; migrate the config to the checkpoint depth before loading." + ) + + +def _build_s3_backend_args(checkpoint_path: str, s3_credential: str | None) -> dict[str, str] | None: + """Build easy_io backend args for an S3 checkpoint when explicit credentials are provided.""" + if not checkpoint_path.startswith("s3://") or s3_credential is None: + return None + return { + "backend": "s3", + "s3_credential_path": s3_credential, + } + + +def _safe_checkpoint_globals() -> list[Any]: + """Return trusted non-tensor classes present in tokenizer training checkpoints.""" + return [ + Any, + AnyNode, + ContainerMetadata, + DictConfig, + ListConfig, + Metadata, + defaultdict, + dict, + int, + list, + ] + + +def load_torch_checkpoint( + checkpoint_path: str | Path, + *, + map_location: TorchMapLocation = "cpu", + s3_credential: str | None = None, + backend_args: dict[str, Any] | None = None, + backend_key: str | None = None, +) -> Any: + """Load a torch checkpoint through easy_io with safe tensor-only unpickling.""" + checkpoint_path_str = str(checkpoint_path) + resolved_backend_args = ( + backend_args if backend_args is not None else _build_s3_backend_args(checkpoint_path_str, s3_credential) + ) + with torch.serialization.safe_globals(_safe_checkpoint_globals()): + return easy_io.load( + checkpoint_path_str, + backend_args=resolved_backend_args, + backend_key=backend_key, + map_location=map_location, + weights_only=True, + ) + + +def load_torch_checkpoint_from_bytes( + checkpoint_bytes: bytes, + *, + map_location: TorchMapLocation = "cpu", +) -> Any: + """Load a torch checkpoint from bytes with safe tensor-only unpickling.""" + with io.BytesIO(checkpoint_bytes) as checkpoint_buffer: + with torch.serialization.safe_globals(_safe_checkpoint_globals()): + return torch.load(checkpoint_buffer, map_location=map_location, weights_only=True) + + +def load_torch_checkpoint_from_easy_io_backend( + backend: Any, + checkpoint_path: str, + *, + map_location: TorchMapLocation = "cpu", +) -> Any: + """Load one checkpoint through an already-configured easy_io backend.""" + return load_torch_checkpoint_from_bytes( + backend.get(filepath=checkpoint_path), + map_location=map_location, + ) + + +def normalize_dcp_model_checkpoint_path(checkpoint_path: str | Path) -> str: + """Map a trainer iteration directory or URI to its DCP model component.""" + checkpoint_path_str = str(checkpoint_path).rstrip("/") + model_component_path = os.path.join(checkpoint_path_str, "model") + if os.path.isdir(model_component_path): + return model_component_path + if checkpoint_path_str.startswith("s3://") and re.search(r"/iter_\d{9}$", checkpoint_path_str): + return model_component_path + return checkpoint_path_str + + +def _cpu_state_placeholder(value: Any) -> Any: + """Create a CPU DCP target without duplicating source-device storage.""" + if isinstance(value, torch.Tensor): + return torch.empty_like(value, device="cpu") # [...] + return value + + +def _add_dcp_ema_template_tensors( + model: nn.Module, + model_state: dict[str, Any], + metadata_keys: set[str], +) -> None: + """Request EMA buffers from DCP when checkpoint metadata says they exist.""" + from cosmos_framework.utils.ema import get_buffer_name + + for name, param in model.named_parameters(): + ema_key = f"ema.{get_buffer_name(name)}" + if ema_key not in metadata_keys or ema_key in model_state: + continue + model_state[ema_key] = torch.empty_like(param.detach(), dtype=torch.float32, device="cpu") # [...] + + +def load_dcp_model_checkpoint( + model: nn.Module, + checkpoint_path: str | Path, + *, + include_ema: bool = False, + allow_partial: bool = False, + checkpoint_key_prefix: str = "", + s3_credential: str | None = None, +) -> dict[str, Any]: + """Load model state from a local or S3 torch.distributed.checkpoint component.""" + from torch.distributed.checkpoint import load + from torch.distributed.checkpoint.filesystem import FileSystemReader + + checkpoint_path_str = normalize_dcp_model_checkpoint_path(checkpoint_path) + + def _build_storage_reader() -> Any: + if checkpoint_path_str.startswith("s3://"): + if not s3_credential: + raise ValueError("s3_credential is required to load a remote DCP checkpoint.") + from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader + + return S3StorageReader( + credential_path=s3_credential, + path=checkpoint_path_str, + ) + return FileSystemReader(checkpoint_path_str) + + storage_reader = _build_storage_reader() + state_dict_metadata = storage_reader.read_metadata().state_dict_metadata + metadata_keys = set(state_dict_metadata) + if include_ema and checkpoint_key_prefix: + raise ValueError("checkpoint_key_prefix cannot be combined with include_ema.") + + current_model_state = model.state_dict() + checkpoint_key_by_model_key = { + model_key: f"{checkpoint_key_prefix}{model_key}" for model_key in current_model_state + } + expected_checkpoint_keys = set(checkpoint_key_by_model_key.values()) + wrapped_matches = sum(f"model.{key}" in metadata_keys for key in expected_checkpoint_keys) + raw_matches = sum(key in metadata_keys for key in expected_checkpoint_keys) + if wrapped_matches == 0 and raw_matches == 0: + raise ValueError( + f"DCP checkpoint {checkpoint_path_str} does not contain tensors matching the requested model " + f"(checkpoint sample: {sorted(metadata_keys)[:5]}; model sample: {sorted(expected_checkpoint_keys)[:5]})." + ) + + metadata_prefix = "model." if wrapped_matches >= raw_matches else "" + component_metadata = { + key.removeprefix(metadata_prefix): value + for key, value in state_dict_metadata.items() + if not metadata_prefix or key.startswith(metadata_prefix) + } + validate_rq_checkpoint_depth( + model, + component_metadata, + checkpoint_path=checkpoint_path_str, + model_keys=current_model_state, + ) + checkpoint_model_state: dict[str, Any] = {} + loaded_model_keys: set[str] = set() + missing_model_keys: set[str] = set() + shape_mismatches: list[tuple[str, tuple[int, ...], tuple[int, ...]]] = [] + for model_key, current_value in current_model_state.items(): + checkpoint_key = checkpoint_key_by_model_key[model_key] + checkpoint_metadata = component_metadata.get(checkpoint_key) + if checkpoint_metadata is None: + missing_model_keys.add(model_key) + continue + checkpoint_shape = getattr(checkpoint_metadata, "size", None) + if isinstance(current_value, torch.Tensor) and checkpoint_shape is not None: + current_shape = tuple(current_value.shape) + saved_shape = tuple(checkpoint_shape) + if current_shape != saved_shape: + shape_mismatches.append((model_key, saved_shape, current_shape)) + continue + checkpoint_model_state[checkpoint_key] = _cpu_state_placeholder(current_value) + loaded_model_keys.add(model_key) + + if shape_mismatches: + raise ValueError( + f"DCP checkpoint {checkpoint_path_str} has {len(shape_mismatches)} model shape mismatches " + f"(sample: {shape_mismatches[:5]})." + ) + if missing_model_keys and not allow_partial: + raise ValueError( + f"DCP checkpoint {checkpoint_path_str} is missing {len(missing_model_keys)} model keys " + f"(sample: {sorted(missing_model_keys)[:10]})." + ) + + if include_ema: + _add_dcp_ema_template_tensors( + model, + checkpoint_model_state, + set(component_metadata), + ) + requested_checkpoint_keys = set(checkpoint_model_state) + unexpected_checkpoint_keys = set(component_metadata) - requested_checkpoint_keys + checkpoint_state = {"model": checkpoint_model_state} if metadata_prefix else checkpoint_model_state + load(checkpoint_state, storage_reader=storage_reader, no_dist=True) + loaded_model_state = { + model_key: checkpoint_model_state[checkpoint_key_by_model_key[model_key]] for model_key in loaded_model_keys + } + loaded_model_state.update( + { + key: value + for key, value in checkpoint_model_state.items() + if key.startswith("ema.") and key not in loaded_model_state + } + ) + return { + "model": loaded_model_state, + DCP_MODEL_LOAD_INFO_KEY: DCPModelLoadInfo( + loaded_model_keys=frozenset(loaded_model_keys), + missing_model_keys=frozenset(missing_model_keys), + unexpected_checkpoint_keys=frozenset(unexpected_checkpoint_keys), + ), + } diff --git a/cosmos_framework/model/tokenizer/evaluation/metric.py b/cosmos_framework/model/tokenizer/evaluation/metric.py index 955b6582..6659dde6 100644 --- a/cosmos_framework/model/tokenizer/evaluation/metric.py +++ b/cosmos_framework/model/tokenizer/evaluation/metric.py @@ -18,6 +18,8 @@ import torch from loguru import logger as logging +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs + def compute_psnr( original: torch.Tensor, @@ -157,7 +159,7 @@ def compute_codebook_usage( all_indices = [] for i, g in enumerate(gathered): all_indices.append(g[: sizes[i].item()]) - flat_indices = torch.cat(all_indices) + flat_indices = cat_with_bounded_inputs(all_indices, dim=0) # Compute code histogram histogram = torch.bincount(flat_indices, minlength=num_codes).float() @@ -318,34 +320,70 @@ def reset(self) -> None: if self._metric is not None: self._metric.reset() + @staticmethod + def _flatten_video_batch(images: torch.Tensor, layout: str = "auto") -> torch.Tensor: + """Flatten image/video tensors into a 4D image batch. + + Args: + images: Tensor shaped [B, C, H, W], [B, C, T, H, W], or [B, T, C, H, W]. + layout: Layout for 5D tensors. Use "bcthw", "btchw", or "auto". + + Returns: + Tensor shaped [B*, C, H, W]. + """ + if images.ndim == 4: + return images + if images.ndim != 5: + raise ValueError(f"FIDComputer.update expected a 4D or 5D tensor, got shape {tuple(images.shape)}") + + normalized_layout = layout.lower() + if normalized_layout == "auto": + channels_at_dim1 = images.shape[1] in (1, 3) + channels_at_dim2 = images.shape[2] in (1, 3) + if channels_at_dim1 == channels_at_dim2: + raise ValueError( + "Ambiguous 5D FID image layout. Pass layout='bcthw' for [B, C, T, H, W] " + "or layout='btchw' for [B, T, C, H, W]." + ) + normalized_layout = "bcthw" if channels_at_dim1 else "btchw" + + if normalized_layout == "bcthw": + images = images.permute(0, 2, 1, 3, 4).contiguous() # [B, T, C, H, W] + elif normalized_layout == "btchw": + images = images.contiguous() # [B, T, C, H, W] + else: + raise ValueError(f"Unsupported FID video layout '{layout}'. Expected 'auto', 'bcthw', or 'btchw'.") + + return images.reshape(-1, *images.shape[2:]) # [B*T, C, H, W] + @torch.no_grad() def update( self, images: torch.Tensor, real: bool = True, + layout: str = "auto", ) -> None: """Update with a batch of images. Separate calls for real and fake images. Args: - images: Images in [0, 1] range, shape (B, C, H, W). + images: Images in [0, 1] range, shaped [B, C, H, W], [B, C, T, H, W], or [B, T, C, H, W]. real: Whether these are real images (True) or fake/reconstructed (False). + layout: Layout for 5D tensors. Use "bcthw", "btchw", or "auto". """ if not self._ensure_initialized(): return - # Handle video tensors (B, C, T, H, W) -> (B*T, C, H, W) - if images.ndim == 5: - images = images.reshape(-1, *images.shape[-3:]) + images = self._flatten_video_batch(images, layout=layout) # [B*, C, H, W] if self._normalize and images.dtype == torch.uint8: - images = images.float() / 255.0 + images = images.float() / 255.0 # [B*, C, H, W] if self._normalize: - images = images.float().to(self.device) + images = images.float().to(self.device) # [B*, C, H, W] else: - images = images.to(self.device) + images = images.to(self.device) # [B*, C, H, W] # Update with autocast disabled for numerical stability with self._autocast_context(): diff --git a/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py b/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py index ab5afa89..732c20eb 100644 --- a/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py +++ b/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py @@ -398,95 +398,6 @@ def calculate_psnr( return psnr.mean() -class Rank0FIDMetric(nn.Module): - """FID metric that runs only on rank 0 to avoid distributed sync issues. - - Uses torchmetrics FrechetInceptionDistance internally but only computes - on rank 0's data to avoid NCCL collective operation mismatches caused by - torchmetrics/torch-fidelity's internal distributed synchronization. - - Note: FID is computed only on rank 0's portion of the data (1/world_size), - which may be less representative than full dataset FID, but avoids - distributed synchronization issues. - - Usage: - fid = Rank0FIDMetric(rank=rank).to(device) - - # During evaluation loop (only rank 0 updates) - for batch in dataloader: - fid.update(real_images, fake_images) - - # Compute FID (only rank 0 has valid result) - if rank == 0: - fid_value = fid.compute() - """ - - def __init__(self, rank: int = 0, feature_dim: int = 2048) -> None: - super().__init__() - self.rank = rank - self.feature_dim = feature_dim - self._fid_metric = None - - # Only initialize FID metric on rank 0 - if self.rank == 0: - try: - from torchmetrics.image.fid import FrechetInceptionDistance - - # normalize=True means input is [0, 1] float, not uint8 - self._fid_metric = FrechetInceptionDistance( - feature=feature_dim, - normalize=True, - sync_on_compute=False, - dist_sync_on_step=False, - ) - except ImportError: - pass - - @torch.no_grad() - def update(self, real_images: torch.Tensor, fake_images: torch.Tensor) -> None: - """Update FID statistics with a batch of real and fake images. - - Only updates on rank 0. - - Args: - real_images: Real images in [0, 1] range, shape (B, C, H, W) or (B, T, C, H, W) - fake_images: Fake/reconstructed images in [0, 1] range - """ - if self.rank != 0 or self._fid_metric is None: - return - - # Handle video format by flattening batch and time dimensions - if real_images.dim() == 5: # (B, T, C, H, W) - real_images = real_images.reshape(-1, *real_images.shape[2:]) - fake_images = fake_images.reshape(-1, *fake_images.shape[2:]) - - # Move metric to same device as images - device = real_images.device - self._fid_metric = self._fid_metric.to(device) - - # torchmetrics FID update - self._fid_metric.update(real_images, real=True) - self._fid_metric.update(fake_images, real=False) - - def compute(self) -> torch.Tensor: - """Compute FID from accumulated statistics. - - Only valid on rank 0. - - Returns: - FID value as a scalar tensor (inf if not rank 0 or metric unavailable) - """ - if self.rank != 0 or self._fid_metric is None: - return torch.tensor(float("inf")) - - return self._fid_metric.compute() - - def reset(self) -> None: - """Reset accumulated statistics.""" - if self._fid_metric is not None: - self._fid_metric.reset() - - __all__ = [ @@ -494,6 +405,5 @@ def reset(self) -> None: "PSNRMetric", "SSIMMetric", "LPIPSMetric", - "Rank0FIDMetric", "calculate_psnr", ] diff --git a/cosmos_framework/model/tokenizer/models/architecture.py b/cosmos_framework/model/tokenizer/models/architecture.py new file mode 100644 index 00000000..f70542b6 --- /dev/null +++ b/cosmos_framework/model/tokenizer/models/architecture.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Shared UniAE architecture specifications used by training and inference.""" + +from __future__ import annotations + +import copy +from typing import Any + +_SIGLIP2_SO400M_COMMON_ARCH: dict[str, Any] = { + "patch_size": (4, 16, 16), + "in_channels": 3072, + "out_channels": 3072, + "encoder_model_channels": 1152, + "encoder_num_blocks": 27, + "encoder_num_heads": 16, + "encoder_mlp_channels": 4304, + "encoder_pe_mode": "joint", + "encoder_qk_rms_norm": False, + "encoder_use_bias": True, + "encoder_use_rms_norm": False, + "decoder_model_channels": 1152, + "decoder_num_blocks": 27, + "decoder_num_heads": 16, + "decoder_mlp_channels": 4304, + "decoder_pe_mode": "joint", + "decoder_qk_rms_norm": True, + "decoder_use_bias": False, + "decoder_use_rms_norm": True, + "use_decoder": True, + "quantizer_type": "rq", + "quantizer_codebook_size": 65536, + # Production checkpoints use four residual-quantization stages. + "quantizer_num_codebooks": 4, + "quantizer_chunk_size": 1, + "use_vf_loss": False, + "freeze_encoder": False, + "pretrained_model_name": "google/siglip2-so400m-patch16-naflex", + "concat_latent": None, + "random_num_sample_frames_batch_sizes": [8, 12, 16, 20, 24], + "inference_num_sample_frames_batch_size": 16, + "inference_num_sample_frames_stride": 16, + "inference_kv_cache_size": 0, +} + + +def get_siglip2_so400m_common_arch() -> dict[str, Any]: + """Return an isolated copy of the shared SigLIP2-SO400M UniAE architecture.""" + return copy.deepcopy(_SIGLIP2_SO400M_COMMON_ARCH) + + +__all__ = ["get_siglip2_so400m_common_arch"] diff --git a/cosmos_framework/model/tokenizer/models/dense_runtime.py b/cosmos_framework/model/tokenizer/models/dense_runtime.py index 93820284..57c08b2d 100644 --- a/cosmos_framework/model/tokenizer/models/dense_runtime.py +++ b/cosmos_framework/model/tokenizer/models/dense_runtime.py @@ -22,6 +22,7 @@ ) from cosmos_framework.model.tokenizer.models.modules.transformer.blocks import LearnedPositionEmbedder from cosmos_framework.model.tokenizer.models.sparse_autoencoder import AutoencoderKL, SparseTransformerBase +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs @dataclass(frozen=True) @@ -50,55 +51,6 @@ class DenseGridMetadata: DenseGridMetadataKey = tuple[str, int, int, int, int, str, str] -class DenseDiagonalGaussianDistribution: - """Diagonal Gaussian posterior for dense channels-last latent tensors.""" - - def __init__(self, parameters: torch.Tensor, deterministic: bool = False) -> None: - """Initialize the dense posterior from `[mean, logvar]` moments.""" - if parameters.ndim not in (4, 5): - raise ValueError( - "DenseDiagonalGaussianDistribution expects 4D/5D channels-last moments, " - f"got shape {tuple(parameters.shape)}." - ) - self.original_dtype = parameters.dtype - self.parameters = parameters.to(torch.float32) - self.mean, self.logvar = torch.chunk(self.parameters, 2, dim=-1) - self.deterministic = deterministic - self.std = torch.exp(0.5 * self.logvar) - self.var = torch.exp(self.logvar) - - if self.deterministic: - self.var = self.std = torch.zeros_like( - self.mean, - device=self.parameters.device, - dtype=self.parameters.dtype, - ) - - def sample(self) -> torch.Tensor: - """Sample a dense channels-last latent tensor.""" - sample = torch.randn_like(self.mean) - return (self.mean + self.std * sample).to(self.original_dtype) - - def kl(self, other: "DenseDiagonalGaussianDistribution" | None = None) -> torch.Tensor: - """Compute KL divergence per latent token, matching sparse scaling.""" - reduce_dims = (-1,) - if self.deterministic: - num_tokens = math.prod(self.mean.shape[:-1]) - return torch.zeros(num_tokens, device=self.parameters.device, dtype=self.parameters.dtype) - if other is None: - kl = 0.5 * torch.sum(torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, dim=reduce_dims) - else: - kl = 0.5 * torch.sum( - torch.pow(self.mean - other.mean, 2) / other.var - + self.var / other.var - - 1.0 - - self.logvar - + other.logvar, - dim=reduce_dims, - ) - return kl.reshape(-1) - - class DenseAutoencoderRuntime(nn.Module): """Dense frozen-runtime wrapper around an existing sparse autoencoder. @@ -406,7 +358,7 @@ def _encode_padded_chunks(padded_chunks: list[torch.Tensor]) -> list[torch.Tenso encoded = self._encode_video_chunk(padded_chunks[0], pad_to=pad_to) return [_trim_boundary_latents(encoded)] - batched_video = torch.cat(padded_chunks, dim=0) # [B*G,t_pad,H,W,3] + batched_video = cat_with_bounded_inputs(padded_chunks, dim=0) # [B*G,t_pad,H,W,3] encoded = self._encode_video_chunk(batched_video, pad_to=pad_to) # [B*G,T_lat,Hp,Wp,2C] per_video_batch = padded_chunks[0].shape[0] return list(_trim_boundary_latents(encoded).split(per_video_batch, dim=0)) @@ -449,7 +401,7 @@ def _encode_padded_chunks(padded_chunks: list[torch.Tensor]) -> list[torch.Tenso if pending_full_chunks: encoded_chunks.extend(_encode_padded_chunks(pending_full_chunks)) - return torch.cat(encoded_chunks, dim=1) + return cat_with_bounded_inputs(encoded_chunks, dim=1) def decode( self, @@ -511,7 +463,7 @@ def decode( if trim_pixel and not is_image: decoded_chunk = decoded_chunk[:, pad_frames:-pad_frames] decoded_chunks.append(decoded_chunk) - return torch.cat(decoded_chunks, dim=1) + return cat_with_bounded_inputs(decoded_chunks, dim=1) def _metadata_cache_key( self, diff --git a/cosmos_framework/model/tokenizer/models/modules/__init__.py b/cosmos_framework/model/tokenizer/models/modules/__init__.py index aae53fc5..ec173616 100644 --- a/cosmos_framework/model/tokenizer/models/modules/__init__.py +++ b/cosmos_framework/model/tokenizer/models/modules/__init__.py @@ -14,7 +14,6 @@ import importlib import os -from typing import Literal from loguru import logger as logging @@ -51,31 +50,6 @@ def _init_from_env() -> None: _init_from_env() -def set_backend(backend: Literal["pytorch", "spconv", "torchsparse"]) -> None: - """Set the sparse tensor backend. - - Args: - backend: Backend to use. - - "pytorch": Pure PyTorch implementation (default, no external dependencies) - - "spconv": Uses spconv.pytorch.SparseConvTensor - - "torchsparse": Uses torchsparse.SparseTensor - """ - global BACKEND - if backend not in _VALID_BACKENDS: - raise ValueError(f"Invalid backend: {backend}. Must be one of {_VALID_BACKENDS}") - BACKEND = backend - - -def set_debug(debug: bool) -> None: - """Enable or disable debug mode. - - Args: - debug: Whether to enable debug mode. - """ - global DEBUG - DEBUG = debug - - # Lazy loading attribute mapping _ATTRIBUTES = { # SparseTensor and operations diff --git a/cosmos_framework/model/tokenizer/models/modules/attention/full_attn.py b/cosmos_framework/model/tokenizer/models/modules/attention/full_attn.py index 78b077f6..06e59eb4 100644 --- a/cosmos_framework/model/tokenizer/models/modules/attention/full_attn.py +++ b/cosmos_framework/model/tokenizer/models/modules/attention/full_attn.py @@ -17,6 +17,7 @@ from cosmos_framework.model.attention.frontend import attention as i4_attention from cosmos_framework.model.attention.varlen import generate_varlen_parameters +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs if TYPE_CHECKING: from cosmos_framework.model.tokenizer.models.modules.sparse_tensor import SparseTensor @@ -191,12 +192,12 @@ def _pack_sparse_temporal_causal_qkv( return empty_q, empty_k, empty_v, [], [], empty_idx return ( - torch.cat(q_feat_chunks, dim=0), - torch.cat(k_feat_chunks, dim=0), - torch.cat(v_feat_chunks, dim=0), + cat_with_bounded_inputs(q_feat_chunks, dim=0), + cat_with_bounded_inputs(k_feat_chunks, dim=0), + cat_with_bounded_inputs(v_feat_chunks, dim=0), q_seqlen, kv_seqlen, - torch.cat(q_index_chunks, dim=0), + cat_with_bounded_inputs(q_index_chunks, dim=0), ) diff --git a/cosmos_framework/model/tokenizer/models/modules/attention/modules.py b/cosmos_framework/model/tokenizer/models/modules/attention/modules.py index fd18129f..4ebcbc2e 100644 --- a/cosmos_framework/model/tokenizer/models/modules/attention/modules.py +++ b/cosmos_framework/model/tokenizer/models/modules/attention/modules.py @@ -24,6 +24,7 @@ sparse_scaled_dot_product_attention, tensor_varlen_scaled_dot_product_attention, ) +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs if TYPE_CHECKING: from cosmos_framework.model.tokenizer.models.modules.sparse_tensor import SparseTensor @@ -89,7 +90,7 @@ def _manual_segmented_scaled_dot_product_attention( ) if out_chunks: - output = torch.cat(out_chunks, dim=0) # [Tq,H,D] + output = cat_with_bounded_inputs(out_chunks, dim=0) # [Tq,H,D] else: output = q.new_empty((0, q.shape[1], v.shape[-1])) # [0,H,D] captured_attn = attn_chunks if store_attn else None @@ -927,7 +928,7 @@ def _concat_temporal_freqs( if not combined_parts: return current_freqs_cis.new_empty((0,) + current_freqs_cis.shape[1:]) - return torch.cat(combined_parts, dim=0) + return cat_with_bounded_inputs(combined_parts, dim=0) combined_parts = [] current_batch_size = current.shape[0] @@ -966,7 +967,7 @@ def _concat_temporal_freqs( if not combined_parts: return current_freqs_cis.new_empty((0,) + current_freqs_cis.shape[1:]) - return torch.cat(combined_parts, dim=0) + return cat_with_bounded_inputs(combined_parts, dim=0) def _concat_temporal_sparse(self, cached: "SparseTensor", current: "SparseTensor") -> "SparseTensor": """Concatenate cached and current sparse tensors along temporal dimension. @@ -1044,9 +1045,9 @@ def _update_kv_cache( offset += batch_len if cached_coords_parts: - selected_coords = torch.cat(cached_coords_parts, dim=0) - cached_k_feats = torch.cat(cached_k_parts, dim=0) - cached_v_feats = torch.cat(cached_v_parts, dim=0) + selected_coords = cat_with_bounded_inputs(cached_coords_parts, dim=0) + cached_k_feats = cat_with_bounded_inputs(cached_k_parts, dim=0) + cached_v_feats = cat_with_bounded_inputs(cached_v_parts, dim=0) else: selected_coords = torch.empty( (0, k.coords.shape[1]), @@ -1087,7 +1088,7 @@ def _update_kv_cache( kv_cache["cached_v"] = cached_v if k_freqs_cis is not None: if cached_freq_parts: - cached_k_freqs_cis = torch.cat(cached_freq_parts, dim=0) + cached_k_freqs_cis = cat_with_bounded_inputs(cached_freq_parts, dim=0) else: cached_k_freqs_cis = k_freqs_cis.new_empty((0,) + k_freqs_cis.shape[1:]) if kv_cache_detach: diff --git a/cosmos_framework/model/tokenizer/models/modules/sparse_tensor.py b/cosmos_framework/model/tokenizer/models/modules/sparse_tensor.py index 4198dd30..268bb695 100644 --- a/cosmos_framework/model/tokenizer/models/modules/sparse_tensor.py +++ b/cosmos_framework/model/tokenizer/models/modules/sparse_tensor.py @@ -17,6 +17,7 @@ import torch from cosmos_framework.model.tokenizer.models.modules import BACKEND, DEBUG +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs SparseTensorData = None # Lazy import @@ -472,8 +473,8 @@ def _select_contiguous_batch_ranges( ) return SparseTensor( - feats=torch.cat(selected_feats_parts, dim=0), - coords=torch.cat(selected_coords_parts, dim=0), + feats=cat_with_bounded_inputs(selected_feats_parts, dim=0), + coords=cat_with_bounded_inputs(selected_coords_parts, dim=0), shape=torch.Size([self.shape[0]] + list(self.feats.shape[1:])), layout=new_layout, has_special_tokens=has_special_tokens, @@ -906,8 +907,8 @@ def __getitem__(self, idx): # Concatenate all coordinates and features (including empty ones) if len(coords) > 0: - coords = torch.cat(coords, dim=0).contiguous() - feats = torch.cat(feats, dim=0).contiguous() + coords = cat_with_bounded_inputs(coords, dim=0).contiguous() + feats = cat_with_bounded_inputs(feats, dim=0).contiguous() else: # Handle case where idx is empty coords = torch.empty( @@ -1431,8 +1432,8 @@ def concat_temporal_at_batch_start(self, other: "SparseTensor") -> "SparseTensor if not has_cached_tokens: return self - combined_coords = torch.cat(combined_coords_parts, dim=0) - combined_feats = torch.cat(combined_feats_parts, dim=0) + combined_coords = cat_with_bounded_inputs(combined_coords_parts, dim=0) + combined_feats = cat_with_bounded_inputs(combined_feats_parts, dim=0) new_shape = torch.Size([current_batch_size] + list(self.feats.shape[1:])) result = SparseTensor( feats=combined_feats, @@ -1499,8 +1500,8 @@ def concat_temporal_at_batch_start(self, other: "SparseTensor") -> "SparseTensor if not has_cached_tokens: return self - combined_coords = torch.cat(combined_coords_parts, dim=0) - combined_feats = torch.cat(combined_feats_parts, dim=0) + combined_coords = cat_with_bounded_inputs(combined_coords_parts, dim=0) + combined_feats = cat_with_bounded_inputs(combined_feats_parts, dim=0) new_shape = torch.Size([current_batch_size] + list(self.feats.shape[1:])) result = SparseTensor( feats=combined_feats, @@ -1631,14 +1632,14 @@ def sparse_cat(inputs: list[SparseTensor], dim: int = 0) -> SparseTensor: coords.append(input.coords.clone()) coords[-1][:, 0] += start start += input.shape[0] - coords = torch.cat(coords, dim=0) - feats = torch.cat([input.feats for input in inputs], dim=0) + coords = cat_with_bounded_inputs(coords, dim=0) + feats = cat_with_bounded_inputs([input.feats for input in inputs], dim=0) output = SparseTensor( coords=coords, feats=feats, ) else: - feats = torch.cat([input.feats for input in inputs], dim=dim) + feats = cat_with_bounded_inputs([input.feats for input in inputs], dim=dim) output = inputs[0].replace(feats) return output @@ -1735,9 +1736,6 @@ def reconstruct_from_temporal_slices( all_coords.append(restored_slice.coords) all_feats.append(restored_slice.feats) - combined_coords = torch.cat(all_coords, dim=0) - combined_feats = torch.cat(all_feats, dim=0) - if target_coords is not None: fast_path = _reconstruct_in_target_batch_order( restored_slices, @@ -1746,6 +1744,8 @@ def reconstruct_from_temporal_slices( ) if fast_path is not None: return fast_path + combined_coords = cat_with_bounded_inputs(all_coords, dim=0) # [N_available, 5] + combined_feats = cat_with_bounded_inputs(all_feats, dim=0) # [N_available, *F] return _match_target_coordinates( combined_coords, combined_feats, @@ -1753,6 +1753,8 @@ def reconstruct_from_temporal_slices( has_special_tokens=has_special_tokens, ) else: + combined_coords = cat_with_bounded_inputs(all_coords, dim=0) # [N_available, 5] + combined_feats = cat_with_bounded_inputs(all_feats, dim=0) # [N_available, *F] sort_key = (combined_coords[:, 0].long() << 20) + combined_coords[:, 1].long() sorted_indices = torch.argsort(sort_key) @@ -1852,12 +1854,12 @@ def _reconstruct_in_target_batch_order( has_special_tokens=False, ) - combined_coords = torch.cat(combined_coords_parts, dim=0) + combined_coords = cat_with_bounded_inputs(combined_coords_parts, dim=0) # [N_target, 5] if combined_coords.shape != target_coords.shape or not torch.equal(combined_coords, target_coords): return None return SparseTensor( - feats=torch.cat(combined_feats_parts, dim=0), + feats=cat_with_bounded_inputs(combined_feats_parts, dim=0), # [N_target, *F] coords=target_coords, shape=torch.Size([batch_size] + list(restored_slices[0].feats.shape[1:])), layout=new_layout, diff --git a/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py b/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py index f69884b9..40b79fd0 100644 --- a/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py +++ b/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py @@ -58,6 +58,7 @@ sparse_to_batched_tensor, sparse_to_img_list, ) +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs, stack_with_bounded_inputs # ============================================================================= # Helper Functions @@ -146,7 +147,9 @@ def _sparse_tensor_to_dense_batch_tokens(sparse_tensor: SparseTensor) -> tuple[t dense_feats = sparse_tensor.feats.reshape(batch_size, seq_len, *feature_shape) return dense_feats, True - dense_feats = torch.stack([sparse_tensor.feats[batch_slice] for batch_slice in sparse_tensor.layout], dim=0) + dense_feats = stack_with_bounded_inputs( + [sparse_tensor.feats[batch_slice] for batch_slice in sparse_tensor.layout], dim=0 + ) # [B,S,*F] return dense_feats, False @@ -154,7 +157,7 @@ def _dense_batch_tokens_to_flat_features(dense_feats: torch.Tensor, used_reshape """Restore dense `[B, S, D]` features to the sparse flat token order.""" if used_reshape_fast_path: return dense_feats.reshape(-1, *dense_feats.shape[2:]) - return torch.cat(list(dense_feats.unbind(dim=0)), dim=0) + return cat_with_bounded_inputs(dense_feats.unbind(dim=0), dim=0) def _crop_temporal_slices_to_ownership( @@ -1208,6 +1211,7 @@ class AutoencoderKLConfig: text_decoder_model_name: str | None = None # e.g., "Qwen/Qwen3-0.6B" or local path text_decoder_family: str = "qwen3" text_decoder_gradient_checkpointing: bool = True + text_decoder_packed_attention_backend: Literal["sdpa", "natten"] = "sdpa" encoder_use_checkpoint: bool | None = None decoder_use_checkpoint: bool | None = None encoder_dense_train_backend: Literal["disabled", "varlen", "batched", "auto"] = "disabled" @@ -1291,6 +1295,7 @@ def __init__( text_decoder_model_name: str | None = None, text_decoder_family: str = "qwen3", text_decoder_gradient_checkpointing: bool = True, + text_decoder_packed_attention_backend: Literal["sdpa", "natten"] = "sdpa", encoder_use_checkpoint: bool | None = None, decoder_use_checkpoint: bool | None = None, encoder_dense_train_backend: Literal["disabled", "varlen", "batched", "auto"] = "disabled", @@ -1311,7 +1316,7 @@ def __init__( task_random_num_sample_frames_batch_sizes: dict[str, list[int]] | None = None, use_dual_latent: bool = False, use_checkpoint: bool = True, - ): + ) -> None: super().__init__() self.use_checkpoint = use_checkpoint self.encoder_use_checkpoint = use_checkpoint if encoder_use_checkpoint is None else encoder_use_checkpoint @@ -1325,6 +1330,7 @@ def __init__( self.use_post_text_decoder = use_post_text_decoder self.spatial_pool_size = spatial_pool_size self.text_decoder_family = text_decoder_family + self.text_decoder_packed_attention_backend: Literal["sdpa", "natten"] = text_decoder_packed_attention_backend self.decoder_temporal_mode = decoder_temporal_mode self.decoder_temporal_query_latent_steps = decoder_temporal_query_latent_steps self.decoder_temporal_cache_latent_steps = decoder_temporal_cache_latent_steps @@ -1435,7 +1441,7 @@ def __init__( elif self.quantizer_type == "rq": self.quantizer = RQBottleneck( latent_shape=(16, 16, latent_channels), - code_shape=(16, 16, 4), + code_shape=(16, 16, self.quantizer_num_codebooks), n_embed=self.quantizer_codebook_size, decay=0.99, shared_codebook=True, @@ -1534,6 +1540,7 @@ def __init__( image_hidden_size=encoder_model_channels, spatial_pool_size=spatial_pool_size, gradient_checkpointing=text_decoder_gradient_checkpointing, + packed_attention_backend=text_decoder_packed_attention_backend, family_spec=get_text_decoder_family_spec( family=text_decoder_family, model_name=text_decoder_model_name, @@ -1548,7 +1555,7 @@ def decode_text( image_feats: "SparseTensor", image_patch_indices: torch.Tensor, segment_ids: torch.Tensor | None = None, - ): + ) -> tuple[Any, int]: """Decode text from image features using the configured text decoder. Passes encoder output (x_no_proj) through spatial merger + causal LM. @@ -1921,7 +1928,7 @@ def decode( logging.warning(f"Decoded shapes are not the same: {[x.shape for x in decoded_list]}") decoded = decoded_list else: - decoded = torch.stack(decoded_list, dim=0) + decoded = stack_with_bounded_inputs(decoded_list, dim=0) # [B,T,C,H,W] decoded = rearrange(decoded, "b t c h w -> b t h w c") if decoded.shape[1] == 1: decoded = decoded.squeeze(1) diff --git a/cosmos_framework/model/tokenizer/models/text_decoder.py b/cosmos_framework/model/tokenizer/models/text_decoder.py index b9f296c1..3cddc135 100644 --- a/cosmos_framework/model/tokenizer/models/text_decoder.py +++ b/cosmos_framework/model/tokenizer/models/text_decoder.py @@ -23,7 +23,6 @@ """ import os -import re from dataclasses import dataclass from types import SimpleNamespace from typing import Any @@ -33,11 +32,15 @@ import torch.nn.functional as F from loguru import logger as logging +from cosmos_framework.model.attention import attention as i4_attention +from cosmos_framework.model.attention.masks import CausalType +from cosmos_framework.model.attention.natten import NATTEN_SUPPORTED, get_natten_version from cosmos_framework.model.tokenizer.utils.hf import ( load_auto_tokenizer_from_cache, prepare_nemotron_tokenizer_snapshot, resolve_hf_snapshot_path, ) +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs, stack_with_bounded_inputs from cosmos_framework.model.tokenizer.utils.vlm_prompt_format import densevl_add_vision_id_text QWEN3_PAD_TOKEN_ID = 151643 @@ -58,6 +61,9 @@ NEMOTRON_2B_THINK_START_TOKEN = "" NEMOTRON_2B_THINK_END_TOKEN = "" TEXT_DECODER_ATTN_IMPLEMENTATION_ENV = "TOKENIZER_TEXT_DECODER_ATTN_IMPLEMENTATION" +PACKED_ATTENTION_BACKEND_SDPA = "sdpa" +PACKED_ATTENTION_BACKEND_NATTEN = "natten" +PACKED_ATTENTION_BACKENDS = frozenset({PACKED_ATTENTION_BACKEND_SDPA, PACKED_ATTENTION_BACKEND_NATTEN}) VQA_THINKING_MODE_OFF = "off" VQA_THINKING_MODE_ON = "on" VQA_THINKING_MODE_RAW = "raw" @@ -88,6 +94,128 @@ def _append_vqa_reasoning_suffix(question: str, reasoning_suffix: str) -> str: return f"{question.rstrip()}{reasoning_suffix}" +def _rotate_half_for_nemotron(hidden_states: torch.Tensor) -> torch.Tensor: + """Rotate the two RoPE feature halves using Nemotron's convention.""" + midpoint = hidden_states.shape[-1] // 2 + first_half = hidden_states[..., :midpoint] # [B,H,T,D/2] + second_half = hidden_states[..., midpoint:] # [B,H,T,D/2] + rotated = torch.cat((-second_half, first_half), dim=-1) # [B,H,T,D] + return rotated + + +def _apply_nemotron_rotary_embeddings( + query_states: torch.Tensor, + key_states: torch.Tensor, + cosine: torch.Tensor, + sine: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply the remote Nemotron model's RoPE formulation to Q and K.""" + rotated_query = _rotate_half_for_nemotron(query_states) # [B,Hq,T,D] + rotated_key = _rotate_half_for_nemotron(key_states) # [B,Hkv,T,D] + query_states_ = query_states * cosine + rotated_query * sine # [B,Hq,T,D] + key_states_ = key_states * cosine + rotated_key * sine # [B,Hkv,T,D] + return query_states_, key_states_ + + +def _forward_nemotron_natten_attention( + attention_module: nn.Module, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + cumulative_seqlen: torch.Tensor, + max_seqlen: int, + rotary_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> torch.Tensor: + """Run one Nemotron self-attention block with NATTEN causal varlen attention.""" + batch_size, sequence_length, _ = hidden_states.shape + if batch_size != 1: + raise ValueError(f"NATTEN varlen requires packed batch size 1, got {batch_size}.") + + query_states = attention_module.q_proj(hidden_states) # [1,T,Hq*D] + query_states = query_states.view( + batch_size, + sequence_length, + attention_module.num_heads, + attention_module.head_dim, + ).transpose(1, 2) # [1,Hq,T,D] + key_states = attention_module.k_proj(hidden_states) # [1,T,Hkv*D] + key_states = key_states.view( + batch_size, + sequence_length, + attention_module.num_key_value_heads, + attention_module.head_dim, + ).transpose(1, 2) # [1,Hkv,T,D] + value_states = attention_module.v_proj(hidden_states) # [1,T,Hkv*D] + value_states = value_states.view( + batch_size, + sequence_length, + attention_module.num_key_value_heads, + attention_module.head_dim, + ).transpose(1, 2) # [1,Hkv,T,D] + + if rotary_embeddings is None: + normalized_position_ids = position_ids.reshape(batch_size, sequence_length).long() # [1,T] + cosine, sine = attention_module.rotary_emb( + value_states, + normalized_position_ids, + ) # each [1,1,T,D] + else: + cosine, sine = rotary_embeddings + query_states_, key_states_ = _apply_nemotron_rotary_embeddings( + query_states, + key_states, + cosine, + sine, + ) # [1,Hq,T,D], [1,Hkv,T,D] + + query_states_ = query_states_.transpose(1, 2).contiguous() # [1,T,Hq,D] + key_states_ = key_states_.transpose(1, 2).contiguous() # [1,T,Hkv,D] + value_states = value_states.transpose(1, 2).contiguous() # [1,T,Hkv,D] + attention_output = i4_attention( + query=query_states_, + key=key_states_, + value=value_states, + is_causal=True, + causal_type=CausalType.DontCare, + cumulative_seqlen_Q=cumulative_seqlen, + cumulative_seqlen_KV=cumulative_seqlen, + max_seqlen_Q=max_seqlen, + max_seqlen_KV=max_seqlen, + backend=PACKED_ATTENTION_BACKEND_NATTEN, + ) + if not isinstance(attention_output, torch.Tensor): + raise TypeError("NATTEN attention unexpectedly returned log-sum-exp output.") + attention_output = attention_output.reshape(batch_size, sequence_length, attention_module.hidden_size) # [1,T,D] + projected_output = attention_module.o_proj(attention_output) # [1,T,D] + return projected_output + + +def _forward_nemotron_natten_decoder_layer( + decoder_layer: nn.Module, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + cumulative_seqlen: torch.Tensor, + max_seqlen: int, + rotary_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> torch.Tensor: + """Run one remote-code Nemotron layer while replacing only self-attention.""" + residual = hidden_states # [1,T,D] + normalized_hidden_states = decoder_layer.input_layernorm(hidden_states) # [1,T,D] + attention_output = _forward_nemotron_natten_attention( + decoder_layer.self_attn, + normalized_hidden_states, + position_ids, + cumulative_seqlen, + max_seqlen, + rotary_embeddings, + ) # [1,T,D] + hidden_states = residual + attention_output # [1,T,D] + + residual = hidden_states # [1,T,D] + normalized_hidden_states = decoder_layer.post_attention_layernorm(hidden_states) # [1,T,D] + mlp_output = decoder_layer.mlp(normalized_hidden_states) # [1,T,D] + return residual + mlp_output # [1,T,D] + + @dataclass(frozen=True) class TextDecoderFamilySpec: """Static configuration for one supported text decoder family.""" @@ -179,22 +307,6 @@ def _repair_nemotron_rotary_buffers(module: nn.Module) -> int: return repaired_count -def _keep_only_first_image_placeholder( - text: str, - image_placeholder: str = "", -) -> str: - """Keep only the first image placeholder occurrence in a prompt string.""" - first_index = text.find(image_placeholder) - if first_index == -1: - return text - - before_first = text[:first_index] - after_first = text[first_index + len(image_placeholder) :] - escaped_placeholder = re.escape(image_placeholder) - after_first_cleaned = re.sub(escaped_placeholder + r"\s*", "", after_first) - return before_first + image_placeholder + after_first_cleaned - - def infer_text_decoder_family(model_name: str) -> str: """Infer the decoder family from a HuggingFace model name.""" model_name_lower = model_name.lower() @@ -437,8 +549,8 @@ def forward( new_layout.append(slice(current_offset, current_offset + num_merged)) current_offset += num_merged - merged_feats = torch.cat(merged_feats_list, dim=0) - merged_coords = torch.cat(merged_coords_list, dim=0) + merged_feats = cat_with_bounded_inputs(merged_feats_list, dim=0) + merged_coords = cat_with_bounded_inputs(merged_coords_list, dim=0) # MLP projection: merged_hidden_size -> out_hidden_size merged_feats = self.linear_fc2(self.act_fn(self.linear_fc1(merged_feats))) @@ -463,7 +575,8 @@ def __init__( dtype: torch.dtype | None = None, attn_implementation: str = "flash_attention_2", family_spec: TextDecoderFamilySpec | None = None, - ): + packed_attention_backend: str = PACKED_ATTENTION_BACKEND_SDPA, + ) -> None: super().__init__() if dtype is None: @@ -472,8 +585,28 @@ def __init__( from transformers import AutoModelForCausalLM self.spec = family_spec or get_text_decoder_family_spec(model_name=model_name) + if packed_attention_backend not in PACKED_ATTENTION_BACKENDS: + raise ValueError( + f"Unsupported packed_attention_backend={packed_attention_backend!r}; " + f"expected one of {sorted(PACKED_ATTENTION_BACKENDS)}." + ) + if packed_attention_backend == PACKED_ATTENTION_BACKEND_NATTEN and self.spec.family != NEMOTRON_2B_SPEC.family: + raise ValueError("NATTEN packed attention currently supports only the Nemotron 2B text decoder.") + if ( + packed_attention_backend == PACKED_ATTENTION_BACKEND_NATTEN + and torch.cuda.is_available() + and not NATTEN_SUPPORTED + ): + raise RuntimeError("NATTEN packed attention was requested but NATTEN is unavailable or unsupported.") + self.packed_attention_backend: str = packed_attention_backend self._model_name: str = model_name or self.spec.default_model_name resolved_attn_implementation = _resolve_attn_implementation(attn_implementation) + if packed_attention_backend == PACKED_ATTENTION_BACKEND_NATTEN: + # Packed training bypasses the remote decoder attention. Keep all + # non-packed calls on the explicit dependency-free fallback. + resolved_attn_implementation = "eager" + if NATTEN_SUPPORTED: + logging.info(f"Using NATTEN {get_natten_version()} for packed text decoder attention") logging.info(f"Loading text decoder: {self._model_name}") logging.info(f"Using text decoder attention backend: {resolved_attn_implementation}") hf_cache_dir = os.environ.get("HF_HOME") @@ -592,7 +725,8 @@ def __init__( logging.info( f"TextDecoderWrapper ready: family={self.spec.family}, hidden_size={hidden_size}, " f"layers={len(self.text_decoder.model.layers)}, " - f"gradient_checkpointing={gradient_checkpointing_enabled}, use_cache=False" + f"gradient_checkpointing={gradient_checkpointing_enabled}, " + f"packed_attention_backend={self.packed_attention_backend}, use_cache=False" ) def _get_eos_token_ids(self) -> tuple[int, ...]: @@ -636,6 +770,77 @@ def _embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: """Embed token IDs using the model's input embedding module.""" return self.text_decoder.get_input_embeddings()(input_ids) + def _forward_nemotron_natten_from_embeddings( + self, + text_embeds: torch.Tensor, + position_ids: torch.Tensor, + cumulative_seqlen: torch.Tensor, + max_seqlen: int, + ) -> SimpleNamespace: + """Run packed Nemotron layers with NATTEN varlen self-attention.""" + if self.spec.family != NEMOTRON_2B_SPEC.family: + raise ValueError("NATTEN packed attention currently supports only the Nemotron 2B text decoder.") + if text_embeds.shape[0] != 1: + raise ValueError(f"NATTEN varlen requires packed batch size 1, got {text_embeds.shape[0]}.") + if cumulative_seqlen.dtype != torch.int32: + raise ValueError(f"NATTEN cumulative sequence lengths must be int32, got {cumulative_seqlen.dtype}.") + + hidden_states = text_embeds # [1,T,D] + decoder_layers = self.text_decoder.model.layers + if not decoder_layers: + raise ValueError("NATTEN packed attention requires at least one Nemotron decoder layer.") + first_attention = decoder_layers[0].self_attn + normalized_position_ids = position_ids.reshape(hidden_states.shape[0], hidden_states.shape[1]).long() # [1,T] + rotary_reference = hidden_states.new_empty( + ( + hidden_states.shape[0], + first_attention.num_key_value_heads, + 0, + first_attention.head_dim, + ) + ) # [1,Hkv,0,D] + rotary_embeddings = first_attention.rotary_emb( + rotary_reference, + normalized_position_ids, + ) # each [1,1,T,D] + use_manual_gradient_checkpointing = ( + self._manual_gradient_checkpointing_enabled and self.training and hidden_states.requires_grad + ) + for decoder_layer in decoder_layers: + if use_manual_gradient_checkpointing: + + def _layer_forward( + layer_hidden_states: torch.Tensor, + layer_module: nn.Module = decoder_layer, + ) -> torch.Tensor: + return _forward_nemotron_natten_decoder_layer( + layer_module, + layer_hidden_states, + position_ids, + cumulative_seqlen, + max_seqlen, + rotary_embeddings, + ) + + hidden_states = torch.utils.checkpoint.checkpoint( + _layer_forward, + hidden_states, + preserve_rng_state=True, + use_reentrant=False, + ) # [1,T,D] + else: + hidden_states = _forward_nemotron_natten_decoder_layer( + decoder_layer, + hidden_states, + position_ids, + cumulative_seqlen, + max_seqlen, + rotary_embeddings, + ) # [1,T,D] + + hidden_states = self.text_decoder.model.norm(hidden_states) # [1,T,D] + return SimpleNamespace(last_hidden_state=hidden_states, past_key_values=None) + def _forward_from_embeddings( self, text_embeds: torch.Tensor, @@ -644,8 +849,30 @@ def _forward_from_embeddings( cache_position: torch.Tensor | None = None, use_cache: bool = False, past_key_values: Any | None = None, + packed_cumulative_seqlen: torch.Tensor | None = None, + packed_max_seqlen: int | None = None, ) -> SimpleNamespace: """Run the LM from caller-provided token embeddings.""" + has_packed_cumulative_seqlen = packed_cumulative_seqlen is not None + has_packed_max_seqlen = packed_max_seqlen is not None + if has_packed_cumulative_seqlen != has_packed_max_seqlen: + raise ValueError("packed_cumulative_seqlen and packed_max_seqlen must be provided together.") + if has_packed_cumulative_seqlen: + if self.packed_attention_backend != PACKED_ATTENTION_BACKEND_NATTEN: + raise ValueError("Packed varlen metadata was provided without the NATTEN packed attention backend.") + if attention_mask is not None or use_cache or past_key_values is not None: + raise ValueError("NATTEN packed training does not support masks or KV-cache inputs.") + if position_ids is None: + raise ValueError("NATTEN packed attention requires per-token position IDs.") + assert packed_cumulative_seqlen is not None + assert packed_max_seqlen is not None + return self._forward_nemotron_natten_from_embeddings( + text_embeds=text_embeds, + position_ids=position_ids, + cumulative_seqlen=packed_cumulative_seqlen, + max_seqlen=packed_max_seqlen, + ) + if self.spec.supports_inputs_embeds_forward: model_kwargs: dict[str, Any] = { "inputs_embeds": text_embeds, @@ -798,10 +1025,11 @@ def _encode(text: str) -> list[int]: image_placeholder = "" num_visuals = len(image_token_counts_by_visual) normalized_question = str(question).strip() + placeholder_count = normalized_question.count(image_placeholder) if num_visuals == 1: - normalized_question = _keep_only_first_image_placeholder(normalized_question, image_placeholder) + if placeholder_count > 1: + raise ValueError(f"VQA prompt has {placeholder_count} image placeholders for one visual input.") else: - placeholder_count = normalized_question.count(image_placeholder) if placeholder_count == 0: visual_placeholders = "\n".join(image_placeholder for _ in image_token_counts_by_visual) normalized_question = f"{visual_placeholders}\n{normalized_question}".strip() @@ -868,10 +1096,9 @@ def forward( ) -> tuple[torch.Tensor, int]: """Forward pass: inject image features into text and run causal LM. - When segment_ids are provided (packed sequences), the packed [1, S] sequence - is split into per-segment batch elements [num_segments, max_seg_len]. - HuggingFace flash attention with unpadding then produces cu_seqlens per batch - element, giving segment-isolated causal attention. + When segment_ids are provided, SDPA reshapes the packed sequence into a + padded per-segment batch, while NATTEN retains the packed layout and uses + cumulative sequence offsets. Both paths provide segment-isolated causal attention. Position IDs reset per segment for correct RoPE. Args: @@ -879,14 +1106,15 @@ def forward( image_feats_tensor: [N, encoder_dim] raw encoder output features. image_coords: [N, 4 or 5] spatial coordinates. image_patch_indices: [N_pooled] flat indices into [B*S] sequence where - merged image features should be inserted. Padded entries are -1. + merged image features should be inserted. The host collator validates + and removes padding sentinels before device transfer. image_layout: Optional batch layout slices for the image features. segment_ids: [B, S] segment IDs for packed sequences. Values >= 0 indicate valid segments, -1 indicates padding. When provided, enables segment-isolated attention and per-segment position IDs. Returns: - lm_logits: [B, S, vocab_size] next-token prediction logits. + Dense [B, S, vocab_size] logits. num_pooled_tokens: Number of image tokens after spatial merging. """ image_features = image_feats_tensor @@ -917,27 +1145,46 @@ def forward( text_embeds = self._embed_input_ids(input_ids) # [B, S, d] B, S, d = text_embeds.shape - # Zero out vision placeholder positions, then insert real image features - vision_mask = (input_ids != self.vision_token_id).to(dtype=text_embeds.dtype) - text_embeds = text_embeds * vision_mask[:, :, None] - - if num_pooled_tokens > 0 and len(image_patch_indices) > 0: - valid_indices = image_patch_indices[image_patch_indices >= 0] - max_idx = B * S - 1 - valid_indices = valid_indices[valid_indices <= max_idx] - num_to_insert = min(len(valid_indices), num_pooled_tokens) + # The host collator validates exact placeholder/index coverage and removes + # padding sentinels before tensors move to the accelerator. + # Zero out vision placeholder positions, then insert real image features. + vision_mask = (input_ids != self.vision_token_id).to(dtype=text_embeds.dtype) # [B,S] + text_embeds = text_embeds * vision_mask[:, :, None] # [B,S,d] - if num_to_insert > 0: - text_embeds_flat = text_embeds.reshape(B * S, d) - insert_idx = valid_indices[:num_to_insert].to(device=text_embeds.device, dtype=torch.long) - text_embeds_flat[insert_idx] = image_features[:num_to_insert].to(text_embeds.dtype) - text_embeds = text_embeds_flat.reshape(B, S, d) + if image_patch_indices.ndim != 1: + raise ValueError( + f"Image patch indices must be one-dimensional, got shape {tuple(image_patch_indices.shape)}." + ) + if ( + image_patch_indices.dtype == torch.bool + or image_patch_indices.is_floating_point() + or image_patch_indices.is_complex() + ): + raise TypeError(f"Image patch indices must use an integer dtype, got {image_patch_indices.dtype}.") + if image_patch_indices.numel() != num_pooled_tokens: + raise ValueError( + "Image feature/index count mismatch: " + f"got {num_pooled_tokens} pooled features and {image_patch_indices.numel()} indices." + ) + if image_patch_indices.numel() > 0: + insert_idx = image_patch_indices.to(device=input_ids.device, dtype=torch.long) # [N] + text_embeds_flat = text_embeds.reshape(B * S, d) # [B*S,d] + text_insert_idx = insert_idx.to(device=text_embeds.device) # [N] + text_embeds_flat[text_insert_idx] = image_features.to(text_embeds.dtype) # [N,d] + text_embeds = text_embeds_flat.reshape(B, S, d) # [B,S,d] # Segment-aware forward pass for packed sequences if segment_ids is not None: - lm_logits = self._forward_packed(text_embeds, input_ids, segment_ids) + lm_logits = self._forward_packed( + text_embeds, + input_ids, + segment_ids, + ) else: - lm_logits = self._forward_standard(text_embeds, input_ids=input_ids) + lm_logits = self._forward_standard( + text_embeds, + input_ids=input_ids, + ) return lm_logits, num_pooled_tokens @@ -992,9 +1239,9 @@ def _forward_packed( ) -> torch.Tensor: """Segment-isolated forward pass for packed sequences. - Splits packed [B, S] into per-segment batch [num_segments, max_seg_len]. - Flash attention with unpadding naturally isolates batch elements via cu_seqlens. - Position IDs reset per segment for correct RoPE. + SDPA reshapes packed [B, S] into padded per-segment batches. NATTEN keeps + valid tokens contiguous and uses cumulative sequence offsets. Position IDs + reset per segment for correct RoPE in both paths. Args: text_embeds: [B, S, d] embeddings with image features already injected. @@ -1002,14 +1249,14 @@ def _forward_packed( segment_ids: [B, S] segment IDs (>=0 for valid, -1 for padding). Returns: - lm_logits: [B, S, vocab_size] reassembled in original packed layout. + Dense [B, S, vocab_size] logits. """ B, S, d = text_embeds.shape device = text_embeds.device dtype = text_embeds.dtype # Process each batch element (typically B=1 for packed sequences) - all_logits = [] + all_logits: list[torch.Tensor] = [] for b in range(B): seg_ids = segment_ids[b] # [S] embeds = text_embeds[b] # [S, d] @@ -1035,42 +1282,54 @@ def _forward_packed( "Segment IDs must be contiguous within packed text decoder inputs" ) num_segments = seg_lengths.numel() - max_seg_len = int(seg_lengths.max().item()) + + segment_starts = torch.cumsum(seg_lengths, dim=0) - seg_lengths # [R] + segment_positions = torch.arange(valid_indices.numel(), device=device, dtype=torch.long) # [V] + segment_positions = segment_positions - torch.repeat_interleave(segment_starts, seg_lengths) # [V] + + if self.packed_attention_backend == PACKED_ATTENTION_BACKEND_NATTEN: + segment_lengths_int32 = seg_lengths.to(dtype=torch.int32) # [R] + zero_offset = torch.zeros(1, device=device, dtype=torch.int32) # [1] + cumulative_lengths = segment_lengths_int32.cumsum(dim=0, dtype=torch.int32) # [R] + cumulative_seqlen = torch.cat([zero_offset, cumulative_lengths], dim=0) # [R+1] + max_seqlen = int(seg_lengths.max().item()) + packed_embeds = valid_embeds.unsqueeze(0) # [1,V,D] + packed_position_ids = segment_positions.unsqueeze(0) # [1,V] + outputs = self._forward_from_embeddings( + text_embeds=packed_embeds, + attention_mask=None, + position_ids=packed_position_ids, + use_cache=False, + packed_cumulative_seqlen=cumulative_seqlen, + packed_max_seqlen=max_seqlen, + ) + valid_logits = self._lm_head(outputs.last_hidden_state) # [1,V,Vocab] + packed_logits = valid_logits.new_zeros((S, valid_logits.shape[-1])) # [S,Vocab] + packed_logits[valid_indices] = valid_logits[0] # [V,Vocab] + all_logits.append(packed_logits) + continue segment_rows = torch.repeat_interleave( torch.arange(num_segments, device=device, dtype=torch.long), seg_lengths, - ) - segment_starts = torch.cumsum(seg_lengths, dim=0) - seg_lengths - segment_positions = torch.arange(valid_indices.numel(), device=device, dtype=torch.long) - segment_positions = segment_positions - torch.repeat_interleave(segment_starts, seg_lengths) - - # Build per-segment batch: [num_segments, max_seg_len, d] - batch_embeds = torch.zeros(num_segments, max_seg_len, d, device=device, dtype=dtype) - batch_embeds[segment_rows, segment_positions] = valid_embeds - - position_template = torch.arange(max_seg_len, device=device, dtype=torch.long) - batch_mask = (position_template.unsqueeze(0) < seg_lengths.unsqueeze(1)).to(dtype=torch.long) - batch_pos = position_template.unsqueeze(0).expand(num_segments, -1) * batch_mask + ) # [V] + max_seg_len = int(seg_lengths.max().item()) + batch_embeds = torch.zeros(num_segments, max_seg_len, d, device=device, dtype=dtype) # [R,T,D] + batch_embeds[segment_rows, segment_positions] = valid_embeds # [V,D] - # Forward through the text decoder with one batch element per segment, - # which isolates attention across packed examples. - cache_position = position_template + position_template = torch.arange(max_seg_len, device=device, dtype=torch.long) # [T] + batch_mask = (position_template.unsqueeze(0) < seg_lengths.unsqueeze(1)).to(dtype=torch.long) # [R,T] + batch_pos = position_template.unsqueeze(0).expand(num_segments, -1) * batch_mask # [R,T] outputs = self._forward_from_embeddings( text_embeds=batch_embeds, attention_mask=batch_mask, position_ids=batch_pos, - cache_position=cache_position, + cache_position=position_template, use_cache=False, ) - seg_logits = self._lm_head(outputs.last_hidden_state) - # seg_logits: [num_segments, max_seg_len, vocab_size] - - # Reassemble into original packed layout [S, vocab_size] - V = seg_logits.shape[-1] - packed_logits = torch.zeros(S, V, device=device, dtype=seg_logits.dtype) - packed_logits[valid_indices] = seg_logits[segment_rows, segment_positions] - + seg_logits = self._lm_head(outputs.last_hidden_state) # [R,T,Vocab] + packed_logits = seg_logits.new_zeros((S, seg_logits.shape[-1])) # [S,Vocab] + packed_logits[valid_indices] = seg_logits[segment_rows, segment_positions] # [V,Vocab] all_logits.append(packed_logits) # Guard: empty batch (all samples skipped by collate) -> empty logits tensor. @@ -1078,7 +1337,7 @@ def _forward_packed( V = self.lm_config.vocab_size return torch.zeros(0, S, V, device=device, dtype=dtype) # Stack batch: [B, S, vocab_size] - return torch.stack(all_logits, dim=0) + return stack_with_bounded_inputs(all_logits, dim=0) # [B,S,Vocab] def _decode_generation_result( self, @@ -1197,7 +1456,8 @@ def _generate_from_prefix_embeddings( if len(generated_tokens) == 0: return input_ids - return torch.cat([input_ids, torch.cat(generated_tokens, dim=1)], dim=1) + generated_ids = cat_with_bounded_inputs(generated_tokens, dim=1) # [B,T_generated] + return torch.cat([input_ids, generated_ids], dim=1) @torch.no_grad() def generate_caption( @@ -1366,6 +1626,7 @@ def _encode(s: str) -> list[int]: return tok.encode(s, add_special_tokens=False) answer_decode_prefix_ids: list[int] | None = None + answer_suppress_token_ids: tuple[int, ...] = () if self.spec.family == QWEN3_SPEC.family: system_turn = ( [QWEN3_IM_START_TOKEN_ID] @@ -1383,6 +1644,7 @@ def _encode(s: str) -> list[int]: # Empty think block signals Qwen3 non-thinking mode. no_think = [QWEN3_THINK_START_TOKEN_ID] + _encode("\n\n") + [QWEN3_THINK_END_TOKEN_ID] + _encode("\n\n") asst_prefix += no_think + answer_suppress_token_ids = self.spec.suppress_token_ids eos_token_ids = self._get_eos_token_ids() elif self.spec.family == NEMOTRON_2B_SPEC.family: im_start_id = _get_required_token_id(tok, NEMOTRON_2B_IM_START_TOKEN) @@ -1398,6 +1660,7 @@ def _encode(s: str) -> list[int]: asst_prefix = [im_start_id] + _encode("assistant\n") if resolved_thinking_mode == VQA_THINKING_MODE_OFF: asst_prefix += [think_id, end_think_id] + answer_suppress_token_ids = (think_id, end_think_id) elif resolved_thinking_mode == VQA_THINKING_MODE_ON: answer_decode_prefix_ids = [think_id] + _encode("\n") asst_prefix += answer_decode_prefix_ids @@ -1436,9 +1699,7 @@ def _encode(s: str) -> list[int]: text_embeds=text_embeds, temperature=temperature, eos_token_ids=eos_token_ids, - suppress_token_ids=( - self.spec.suppress_token_ids if resolved_thinking_mode == VQA_THINKING_MODE_OFF else () - ), + suppress_token_ids=answer_suppress_token_ids, ) answer, metadata = self._decode_generation_result( diff --git a/cosmos_framework/model/tokenizer/models/utils.py b/cosmos_framework/model/tokenizer/models/utils.py index d105228f..7792e397 100644 --- a/cosmos_framework/model/tokenizer/models/utils.py +++ b/cosmos_framework/model/tokenizer/models/utils.py @@ -203,18 +203,23 @@ def batch_tensor_to_sparse(batch_tensor: torch.Tensor, patch_size: tuple[int, in """Convert a batch tensor to SparseTensor. Args: - batch_tensor: Input tensor of shape [B, C, H, W] or [B, T, H, W, C]. + batch_tensor: Input tensor of shape [B, C, H, W], [B, H, W, C], or [B, T, H, W, C]. patch_size: Patch size (Pt, Ph, Pw). Returns: SparseTensor with patches as features. """ + Pt, Ph, Pw = patch_size if batch_tensor.dim() == 4: if batch_tensor.shape[1] == 3: - batch_tensor = rearrange(batch_tensor, "b c h w -> b h w c") - batch_tensor = batch_tensor.unsqueeze(1) + batch_tensor = rearrange(batch_tensor, "b c h w -> b h w c") # [B,H,W,C] + batch_tensor = batch_tensor.unsqueeze(1) # [B,1,H,W,C] + if Pt > 1: + temporal_padding = batch_tensor.new_zeros( + (batch_tensor.shape[0], Pt - 1, *batch_tensor.shape[2:]) + ) # [B,Pt-1,H,W,C] + batch_tensor = torch.cat((batch_tensor, temporal_padding), dim=1) # [B,Pt,H,W,C] batch, T, H, W, _ = batch_tensor.shape - Pt, Ph, Pw = patch_size assert T % Pt == 0 and H % Ph == 0 and W % Pw == 0 num_temporal_patches = T // Pt num_height_patches = H // Ph diff --git a/cosmos_framework/model/tokenizer/utils/tensors.py b/cosmos_framework/model/tokenizer/utils/tensors.py new file mode 100644 index 00000000..62810a19 --- /dev/null +++ b/cosmos_framework/model/tokenizer/utils/tensors.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Tensor operation helpers shared by tokenizer training and inference.""" + +from collections.abc import Sequence + +import torch + +# PyTorch's CUDA multi-input cat path issued an illegal memory access on +# Blackwell with 130 and 194 views, while 98 passed. Keep a safety margin. +MAX_CAT_INPUTS_PER_CALL = 64 + + +def _normalize_tensor_sequence( + tensors: Sequence[torch.Tensor], +) -> list[torch.Tensor] | tuple[torch.Tensor, ...]: + """Return a PyTorch-compatible tensor sequence without copying common inputs.""" + if isinstance(tensors, (list, tuple)): + return tensors + return list(tensors) + + +def cat_with_bounded_inputs(tensors: Sequence[torch.Tensor], dim: int = 0) -> torch.Tensor: + """Concatenate tensors without exceeding the CUDA multi-input kernel's safe fan-in.""" + tensor_inputs = _normalize_tensor_sequence(tensors) + if len(tensor_inputs) <= MAX_CAT_INPUTS_PER_CALL: + return torch.cat(tensor_inputs, dim=dim) # [*cat_shape] + + current_level = list(tensor_inputs) + first_device = current_level[0].device + if any(tensor.device != first_device for tensor in current_level[1:]): + devices = ", ".join(sorted({str(tensor.device) for tensor in current_level})) + raise RuntimeError( + "cat_with_bounded_inputs requires tensors to share one device when the input count " + f"exceeds {MAX_CAT_INPUTS_PER_CALL}; got devices: {devices}" + ) + + target_dtype = current_level[0].dtype + for tensor in current_level[1:]: + target_dtype = torch.promote_types(target_dtype, tensor.dtype) + if any(tensor.dtype != target_dtype for tensor in current_level): + current_level = [tensor.to(dtype=target_dtype) for tensor in current_level] # [*input_shapes] + + while len(current_level) > MAX_CAT_INPUTS_PER_CALL: + # Each partial preserves all dimensions except for the concatenated axis. + current_level = [ + torch.cat(current_level[start : start + MAX_CAT_INPUTS_PER_CALL], dim=dim) + for start in range(0, len(current_level), MAX_CAT_INPUTS_PER_CALL) + ] # [*partial_cat_shapes] + return torch.cat(current_level, dim=dim) # [*cat_shape] + + +def stack_with_bounded_inputs(tensors: Sequence[torch.Tensor], dim: int = 0) -> torch.Tensor: + """Stack tensors without exposing the underlying cat kernel to unbounded fan-in.""" + tensor_inputs = _normalize_tensor_sequence(tensors) + if len(tensor_inputs) <= MAX_CAT_INPUTS_PER_CALL: + return torch.stack(tensor_inputs, dim=dim) # [*stack_shape] + expanded_inputs = [tensor.unsqueeze(dim) for tensor in tensor_inputs] # [*expanded_input_shapes] + return cat_with_bounded_inputs(expanded_inputs, dim=dim) # [*stack_shape] + + +__all__ = ["MAX_CAT_INPUTS_PER_CALL", "cat_with_bounded_inputs", "stack_with_bounded_inputs"] diff --git a/cosmos_framework/tools/flops/qwen3_vl.py b/cosmos_framework/tools/flops/qwen3_vl.py index c84a2444..7705795b 100644 --- a/cosmos_framework/tools/flops/qwen3_vl.py +++ b/cosmos_framework/tools/flops/qwen3_vl.py @@ -65,6 +65,7 @@ def compute_attention_flops( num_kv_heads: int, head_dim: int | None = None, is_causal: bool = False, + has_bias: bool = True, ) -> int: """Compute FLOPs for attention mechanism. @@ -81,6 +82,7 @@ def compute_attention_flops( count and the kernel work for the causal case. Defaults to False for bidirectional attention (e.g. the vision encoder); set True for causal decoders. + has_bias: Whether Q/K/V/O projections include bias terms. Returns: Total FLOPs for attention @@ -89,9 +91,9 @@ def compute_attention_flops( head_dim = hidden_size // num_heads # QKV projection: 3 linear layers (but KV uses num_kv_heads) - q_proj_flops = compute_linear_flops(hidden_size, num_heads * head_dim, seq_len, has_bias=True) - k_proj_flops = compute_linear_flops(hidden_size, num_kv_heads * head_dim, seq_len, has_bias=True) - v_proj_flops = compute_linear_flops(hidden_size, num_kv_heads * head_dim, seq_len, has_bias=True) + q_proj_flops = compute_linear_flops(hidden_size, num_heads * head_dim, seq_len, has_bias=has_bias) + k_proj_flops = compute_linear_flops(hidden_size, num_kv_heads * head_dim, seq_len, has_bias=has_bias) + v_proj_flops = compute_linear_flops(hidden_size, num_kv_heads * head_dim, seq_len, has_bias=has_bias) # Causal masking halves the work that scales with S^2 (QK^T, softmax, # attn @ V). The QKV / output projections are not affected. @@ -109,7 +111,7 @@ def compute_attention_flops( attn_v_matmul_flops = int(2 * num_heads * seq_len * seq_len * head_dim * causal_factor) # Output projection - o_proj_flops = compute_linear_flops(num_heads * head_dim, hidden_size, seq_len, has_bias=True) + o_proj_flops = compute_linear_flops(num_heads * head_dim, hidden_size, seq_len, has_bias=has_bias) total_flops = ( q_proj_flops @@ -335,6 +337,7 @@ def compute_text_decoder_flops( num_text_layers: int, head_dim: int = 128, is_causal: bool = True, + attention_bias: bool = False, # MoE parameters num_experts: int = 0, num_experts_per_tok: int = 0, @@ -357,6 +360,7 @@ def compute_text_decoder_flops( causal, so this defaults to True; set to False to recover the bidirectional upper bound (e.g. when comparing against a fitted runtime curve calibrated against the upper bound). + attention_bias: Whether Q/K/V/O attention projections include bias. num_experts: Total number of experts in MoE (0 means no MoE) num_experts_per_tok: Number of experts activated per token (top-k) moe_intermediate_size: Intermediate dimension for each MoE expert @@ -387,6 +391,7 @@ def compute_text_decoder_flops( num_key_value_heads, head_dim, is_causal=is_causal, + has_bias=attention_bias, ) # MLP or MoE (Qwen uses SwiGLU) @@ -409,7 +414,7 @@ def compute_text_decoder_flops( ln_flops = 2 * compute_layernorm_flops(total_tokens, hidden_size) # RMSNorm for Q and K (applied to each head dimension) - qk_norm_flops = 2 * compute_layernorm_flops(total_tokens * num_attention_heads, head_dim) + qk_norm_flops = compute_layernorm_flops(total_tokens * (num_attention_heads + num_key_value_heads), head_dim) decoder_flops += attn_flops + mlp_flops + ln_flops + qk_norm_flops @@ -435,6 +440,7 @@ def compute_qwen3vl_flops( include_lm_head: bool = True, spatial_merge_size: int = 2, is_causal: bool = True, + attention_bias: bool = False, # MoE parameters num_experts: int = 0, num_experts_per_tok: int = 0, @@ -466,6 +472,8 @@ def compute_qwen3vl_flops( is_causal: If True (default), apply the causal-attention factor of 0.5 to the text decoder's S^2 attention terms. The vision encoder is always treated as bidirectional. + attention_bias: Whether text-decoder Q/K/V/O attention projections + include bias. num_experts: Total number of experts in MoE (0 means no MoE) num_experts_per_tok: Number of experts activated per token (top-k) moe_intermediate_size: Intermediate dimension for each MoE expert @@ -540,6 +548,7 @@ def compute_qwen3vl_flops( num_text_layers=num_text_layers, head_dim=head_dim, is_causal=is_causal, + attention_bias=attention_bias, # MoE parameters num_experts=num_experts, num_experts_per_tok=num_experts_per_tok, @@ -612,6 +621,7 @@ def compute_qwen3vl_flops_from_config( """ # Extract MoE parameters if available (MoE models) text_config = config.text_config + attention_bias = getattr(text_config, "attention_bias", False) num_experts = getattr(text_config, "num_experts", 0) num_experts_per_tok = getattr(text_config, "num_experts_per_tok", 0) moe_intermediate_size = getattr(text_config, "moe_intermediate_size", None) @@ -635,6 +645,7 @@ def compute_qwen3vl_flops_from_config( head_dim=text_config.head_dim, spatial_merge_size=config.vision_config.spatial_merge_size, is_causal=is_causal, + attention_bias=attention_bias, # MoE parameters num_experts=num_experts, num_experts_per_tok=num_experts_per_tok, diff --git a/cosmos_framework/tools/flops/test_qwen3_vl.py b/cosmos_framework/tools/flops/test_qwen3_vl.py new file mode 100644 index 00000000..5413598c --- /dev/null +++ b/cosmos_framework/tools/flops/test_qwen3_vl.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from types import SimpleNamespace + +import pytest + +from cosmos_framework.tools.flops.qwen3_vl import ( + compute_attention_flops, + compute_layernorm_flops, + compute_mlp_flops, + compute_qwen3vl_flops, + compute_qwen3vl_flops_from_config, + compute_text_decoder_flops, +) + + +@pytest.mark.L0 +def test_attention_bias_controls_qkvo_projection_bias_terms() -> None: + seq_len = 3 + hidden_size = 8 + num_heads = 2 + num_kv_heads = 1 + head_dim = 4 + + without_bias = compute_attention_flops( + seq_len=seq_len, + hidden_size=hidden_size, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + has_bias=False, + ) + with_bias = compute_attention_flops( + seq_len=seq_len, + hidden_size=hidden_size, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + has_bias=True, + ) + + q_bias_flops = seq_len * num_heads * head_dim + k_bias_flops = seq_len * num_kv_heads * head_dim + v_bias_flops = seq_len * num_kv_heads * head_dim + o_bias_flops = seq_len * hidden_size + assert with_bias - without_bias == q_bias_flops + k_bias_flops + v_bias_flops + o_bias_flops + + +@pytest.mark.L0 +def test_text_decoder_qk_norm_counts_query_and_kv_heads_for_gqa() -> None: + total_tokens = 3 + hidden_size = 8 + intermediate_size = 16 + num_attention_heads = 2 + num_key_value_heads = 1 + head_dim = 4 + + actual = compute_text_decoder_flops( + total_tokens=total_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + num_text_layers=1, + head_dim=head_dim, + is_causal=False, + attention_bias=False, + ) + + attn_flops = compute_attention_flops( + seq_len=total_tokens, + hidden_size=hidden_size, + num_heads=num_attention_heads, + num_kv_heads=num_key_value_heads, + head_dim=head_dim, + is_causal=False, + has_bias=False, + ) + mlp_flops = compute_mlp_flops(total_tokens, hidden_size, intermediate_size, use_swiglu=True) + layernorm_flops = 2 * compute_layernorm_flops(total_tokens, hidden_size) + qk_norm_flops = compute_layernorm_flops( + total_tokens * (num_attention_heads + num_key_value_heads), + head_dim, + ) + + assert actual == attn_flops + mlp_flops + layernorm_flops + qk_norm_flops + + +@pytest.mark.L0 +def test_qwen3vl_flops_from_config_uses_attention_bias() -> None: + text_config = SimpleNamespace( + num_hidden_layers=1, + hidden_size=8, + intermediate_size=16, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=4, + vocab_size=32, + attention_bias=False, + ) + vision_config = SimpleNamespace( + depth=0, + hidden_size=4, + intermediate_size=8, + num_heads=1, + spatial_merge_size=2, + ) + config = SimpleNamespace(text_config=text_config, vision_config=vision_config) + + from_config = compute_qwen3vl_flops_from_config( + config=config, + total_tokens=5, + visual_tokens=1, + num_patches=None, + is_causal=False, + ) + direct_without_bias = compute_qwen3vl_flops( + num_text_layers=text_config.num_hidden_layers, + num_vision_layers=vision_config.depth, + hidden_size=text_config.hidden_size, + intermediate_size=text_config.intermediate_size, + num_attention_heads=text_config.num_attention_heads, + num_key_value_heads=text_config.num_key_value_heads, + vision_hidden_size=vision_config.hidden_size, + vision_intermediate_size=vision_config.intermediate_size, + vision_num_heads=vision_config.num_heads, + vocab_size=text_config.vocab_size, + total_tokens=5, + visual_tokens=1, + num_patches=None, + head_dim=text_config.head_dim, + spatial_merge_size=vision_config.spatial_merge_size, + is_causal=False, + attention_bias=False, + ) + direct_with_bias = compute_qwen3vl_flops( + num_text_layers=text_config.num_hidden_layers, + num_vision_layers=vision_config.depth, + hidden_size=text_config.hidden_size, + intermediate_size=text_config.intermediate_size, + num_attention_heads=text_config.num_attention_heads, + num_key_value_heads=text_config.num_key_value_heads, + vision_hidden_size=vision_config.hidden_size, + vision_intermediate_size=vision_config.intermediate_size, + vision_num_heads=vision_config.num_heads, + vocab_size=text_config.vocab_size, + total_tokens=5, + visual_tokens=1, + num_patches=None, + head_dim=text_config.head_dim, + spatial_merge_size=vision_config.spatial_merge_size, + is_causal=False, + attention_bias=True, + ) + + assert from_config == direct_without_bias + assert from_config["total_flops"] < direct_with_bias["total_flops"] diff --git a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py index fdc944f8..c417cf59 100644 --- a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py +++ b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py @@ -145,6 +145,18 @@ def dump_to_fileobj( obj = obj.cpu().numpy() h, w = obj.shape[1:-1] + # H.264 + yuv420p requires both dimensions to be even. Some eval videos + # keep the raw camera aspect ratio (for example 569x640), and ffmpeg + # fails with "width not divisible by 2" unless we pad before encoding. + # Pad minimally on the bottom/right edge so content coordinates remain + # unchanged and update ffmpeg's raw input size to match the padded array. + pad_h = h % 2 + pad_w = w % 2 + if pad_h or pad_w: + log.debug(f"Padding video from {w}x{h} to {w + pad_w}x{h + pad_h} for MP4 encoding") + obj = np.pad(obj, ((0, 0), (0, pad_h), (0, pad_w), (0, 0)), mode="edge") + h, w = obj.shape[1:-1] + # Default ffmpeg params that ensure width and height are set default_ffmpeg_params = ["-s", f"{w}x{h}"] @@ -160,13 +172,13 @@ def dump_to_fileobj( "output_params": ["-f", "mp4"], } else: - # Use provided ffmpeg_params if any, otherwise use defaults - final_ffmpeg_params = ffmpeg_params if ffmpeg_params is not None else default_ffmpeg_params + # Preserve caller-provided ffmpeg params, but always append the post-padding + # input size so custom params cannot leave ffmpeg seeing stale odd dimensions. mimsave_kwargs = { "fps": fps, "quality": quality, "macro_block_size": 1, - "ffmpeg_params": final_ffmpeg_params, + "ffmpeg_params": (ffmpeg_params or []) + default_ffmpeg_params, "output_params": ["-f", "mp4"], } # Update with any other kwargs diff --git a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler_test.py b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler_test.py new file mode 100644 index 00000000..121ea02e --- /dev/null +++ b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler_test.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from io import BytesIO + +import numpy as np +import pytest + +from cosmos_framework.utils.easy_io.handlers.imageio_video_handler import ImageioVideoHandler + + +@pytest.mark.L0 +def test_dump_to_fileobj_pads_odd_dimensions_for_mp4(monkeypatch): + calls = {} + + def fake_mimsave(file, obj, format, **kwargs): # pylint: disable=redefined-builtin + calls["file"] = file + calls["obj"] = obj + calls["format"] = format + calls["kwargs"] = kwargs + + monkeypatch.setattr( + "cosmos_framework.utils.easy_io.handlers.imageio_video_handler.imageio.mimsave", + fake_mimsave, + ) + + video = np.arange(2 * 3 * 5 * 3, dtype=np.uint8).reshape(2, 3, 5, 3) + ImageioVideoHandler().dump_to_fileobj(video, BytesIO(), crf=25) + + saved = calls["obj"] + assert saved.shape == (2, 4, 6, 3) + np.testing.assert_array_equal(saved[:, :3, :5, :], video) + np.testing.assert_array_equal(saved[:, 3:, :5, :], video[:, 2:3, :, :]) + np.testing.assert_array_equal(saved[:, :3, 5:, :], video[:, :, 4:5, :]) + assert calls["kwargs"]["ffmpeg_params"][-2:] == ["-s", "6x4"] + + +@pytest.mark.L0 +def test_dump_to_fileobj_appends_padded_size_to_custom_non_crf_ffmpeg_params(monkeypatch): + calls = {} + + def fake_mimsave(file, obj, format, **kwargs): # pylint: disable=redefined-builtin + calls["file"] = file + calls["obj"] = obj + calls["format"] = format + calls["kwargs"] = kwargs + + monkeypatch.setattr( + "cosmos_framework.utils.easy_io.handlers.imageio_video_handler.imageio.mimsave", + fake_mimsave, + ) + + video = np.zeros((2, 3, 5, 3), dtype=np.uint8) + ImageioVideoHandler().dump_to_fileobj(video, BytesIO(), ffmpeg_params=["-an"]) + + assert calls["obj"].shape == (2, 4, 6, 3) + assert calls["kwargs"]["ffmpeg_params"] == ["-an", "-s", "6x4"] diff --git a/cosmos_framework/utils/generator/aux_optimizer_utils.py b/cosmos_framework/utils/generator/aux_optimizer_utils.py new file mode 100644 index 00000000..c12356d4 --- /dev/null +++ b/cosmos_framework/utils/generator/aux_optimizer_utils.py @@ -0,0 +1,316 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Shared parameter-categorization for the orthogonalizing optimizers (Muon / Dion2). + +Both ``MuonWithAuxAdamW`` and ``Dion2WithAuxAdamW`` only apply their +orthogonalized (Newton-Schulz) update to *hidden* ``nn.Linear`` weight matrices. +Everything else -- token / positional embeddings, the output head (``lm_head``), +layer-norm scales, biases, and any non-Linear parameter -- must use the auxiliary +AdamW. + +The tricky part is reliably identifying the **embeddings** and the **output head** +across every model architecture in the repo (Qwen3, GPT-OSS, DeepSeekV3, Gemma4, +Qwen3-VL-MoE, unified MoT, ...). We do not rely on a single hard-coded name. +Instead :func:`split_orthogonalizable_params` combines four signals: + +1. **Module type** -- ``nn.Embedding`` weights always go to AdamW (and they are + never ``nn.Linear`` so they would never reach Muon anyway). +2. **Name keywords** -- a configurable set of substrings (default ``{"lm_head"}``, + the convention used by every LLM in this repo) marks output-head Linears. +3. **Tied weights** -- a Linear whose ``weight`` tensor is *the same object* as an + ``nn.Embedding`` weight (``tie_word_embeddings=True``) is the tied output head. +4. **Vocabulary shape** -- a Linear whose output dimension equals some + ``nn.Embedding.num_embeddings`` (the vocab size) is treated as an output + projection. + +Signals (2)-(4) are OR-ed, so an output head is excluded from Muon/Dion2 even if a +new architecture names it differently or ties it to the embedding. The failure +mode of the heuristics is conservative: a misclassified weight falls back to +AdamW (always safe) rather than being orthogonalized. +""" + +import torch +import torch.nn as nn + +# Substrings (matched against the dotted module name) that force an ``nn.Linear`` +# onto the AdamW side. Every LLM in the repo names its head ``lm_head``; the extra +# entries cover common alternative names used elsewhere. +DEFAULT_ADAMW_MODULE_KEYWORDS: tuple[str, ...] = ("lm_head", "embed_out", "output_layer") + + +def is_output_head_linear( + module_name: str, + *, + adamw_module_keywords: tuple[str, ...], +) -> bool: + """Whether an ``nn.Linear`` is an output head (and must use AdamW, not Muon). + + Name-based: every LLM backbone in the repo names its head with an + ``adamw_module_keywords`` substring (default ``lm_head``). Embeddings are + excluded separately by module type in :func:`split_orthogonalizable_params` + (``nn.Embedding`` never reaches this function), and a tied head shares its + weight object with an embedding, so it is caught there too -- hence the name + check alone is sufficient for every current model. + """ + return any(keyword in module_name for keyword in adamw_module_keywords) + + +def split_orthogonalizable_params( + model: nn.Module, + optimizer_param_ids: set[int], + adamw_module_keywords: tuple[str, ...] = DEFAULT_ADAMW_MODULE_KEYWORDS, + expert_param_keywords: tuple[str, ...] = (), +) -> tuple[list[nn.Parameter], list[nn.Parameter], dict[nn.Parameter, str]]: + """Split ``model``'s trainable params by whether they can be orthogonalized. + + Args: + model: The (sub)module to walk -- typically the trainable ``net``. + optimizer_param_ids: ``id()`` of every parameter actually owned by the + optimizer's param groups. Only these are categorized, so that + ``keys_to_select`` filtering is respected and ``state_dict`` stays + consistent. + adamw_module_keywords: Name substrings that force an ``nn.Linear`` onto the + unorthogonalizable (AdamW) side (output heads). + expert_param_keywords: Name substrings that mark **stacked MoE expert** + parameters -- raw ``nn.Parameter`` tensors of shape ``[num_experts, M, + N]`` (e.g. ``gate_up_proj`` / ``down_proj``). When a 3-D+ param matches, + it is routed to the orthogonalizable side so the optimizer can treat + each expert slice as its own matrix. Empty (default) keeps expert + params on AdamW, i.e. no behavior change. + + NOTE (sharding assumption): the optimizer orthogonalizes these per + expert slice assuming the tensor is sharded on dim 0 (the expert axis), + which holds for the FSDP2 ``fully_shard`` path used here (FSDP2 shards + every parameter on dim 0) and makes the update communication-free. It is + NOT guaranteed if tensor/expert parallelism shards *within* an expert + matrix (dim 1/2); that case is unsupported and rejected at step time by + ``MuonWithAuxAdamW._step_stacked_muon`` / ``Dion2WithAuxAdamW._step_stacked_dion2``. + This was not exhaustively audited across every parallelization config, + hence the runtime guard there. + + TODO(expert-parallelism): support 2-D sharding of stacked expert params. + With expert parallelism, dim 0 would be sharded along the expert axis and + dim 1 along the FSDP axis simultaneously, so each rank holds a shard of a + *slice* of each expert matrix rather than whole expert matrices. The + per-expert-slice orthogonalization here assumes dim-0-only sharding (whole + expert matrices are local), so it would need to reconstruct each expert + matrix across the FSDP axis (an all-gather along dim 1, like the dense + 2-D path) before running Newton-Schulz. Doable, but needs changes. + + Returns: + ``(orthogonalizable, unorthogonalizable, param_to_name)`` where + ``orthogonalizable`` holds the hidden ``nn.Linear`` weights (2-D) and, when + ``expert_param_keywords`` is set, the stacked expert tensors (3-D); the + optimizer splits those by rank. ``unorthogonalizable`` is everything else + (embeddings, output head, biases, norms, other non-Linear params), which + uses the auxiliary AdamW. + """ + orthogonalizable: list[nn.Parameter] = [] + unorthogonalizable: list[nn.Parameter] = [] + param_to_name: dict[nn.Parameter, str] = {} + categorized: set[int] = set() + + for name, param in model.named_parameters(): + if param.requires_grad: + param_to_name[param] = name + + def _eligible(p: nn.Parameter) -> bool: + return p.requires_grad and id(p) in optimizer_param_ids and id(p) not in categorized + + def _is_stacked_expert(param: nn.Parameter) -> bool: + if not expert_param_keywords or param.ndim < 3: + return False + name = param_to_name.get(param, "") + return any(keyword in name for keyword in expert_param_keywords) + + for module_name, module in model.named_modules(): + if isinstance(module, nn.Linear): + is_head = is_output_head_linear( + module_name, + adamw_module_keywords=adamw_module_keywords, + ) + if _eligible(module.weight): + (unorthogonalizable if is_head else orthogonalizable).append(module.weight) + categorized.add(id(module.weight)) + if module.bias is not None and _eligible(module.bias): + unorthogonalizable.append(module.bias) + categorized.add(id(module.bias)) + else: + # Embeddings (nn.Embedding), norms, conv, and any raw nn.Parameter -> + # unorthogonalizable (AdamW), except stacked MoE expert tensors when + # opted in. ``recurse=False`` so each owning module handles its own + # params exactly once. + for param in module.parameters(recurse=False): + if not _eligible(param): + continue + (orthogonalizable if _is_stacked_expert(param) else unorthogonalizable).append(param) + categorized.add(id(param)) + + return orthogonalizable, unorthogonalizable, param_to_name + + +# ----------------------------------------------------------------------------- +# Shared orthogonalization / momentum math (used by both MuonWithAuxAdamW and +# Dion2WithAuxAdamW). Kept here so the two optimizers do not duplicate them. +# ----------------------------------------------------------------------------- + + +@torch.compile +def zeropower_via_newtonschulz5(G: torch.Tensor, steps: int = 5) -> torch.Tensor: + """ + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. + + Uses a quintic iteration whose coefficients are selected to maximize the slope at zero. + This produces US'V^T where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which + empirically does not hurt model performance relative to exact UV^T. + + Args: + G: Input tensor of shape (m, n) where m, n >= 1. + steps: Number of Newton-Schulz iterations. + + Returns: + Orthogonalized tensor of same shape as G. + """ + assert G.ndim == 2, "Input must be 2D" + a, b, c = (3.4445, -4.7750, 2.0315) + X = G.bfloat16() + + # Transpose if tall matrix for numerical stability + if G.size(0) > G.size(1): + X = X.T + + # Ensure spectral norm is at most 1 (global norm, matching Moonlight) + X = X / (X.norm() + 1e-7) + + # Perform Newton-Schulz iterations + for _ in range(steps): + A = X @ X.T + B = b * A + c * A @ A + X = a * X + B @ X + + # Transpose back if we transposed earlier + if G.size(0) > G.size(1): + X = X.T + + return X + + +@torch.compile +def zeropower_via_newtonschulz5_batched(G: torch.Tensor, steps: int = 5) -> torch.Tensor: + """Batched Newton-Schulz over a stack of matrices. + + Same quintic iteration as :func:`zeropower_via_newtonschulz5`, but applied to a + batch ``G`` of shape ``[..., M, N]`` (e.g. stacked MoE experts ``[E, M, N]``) + using batched matmuls. All matrices in the batch share ``M, N``, so the + transpose-if-tall decision is uniform across the batch. + """ + assert G.ndim >= 3, "Batched input must be at least 3D ([..., M, N])" + a, b, c = (3.4445, -4.7750, 2.0315) + X = G.bfloat16() + + transposed = False + if X.size(-2) > X.size(-1): + X = X.mT + transposed = True + + # Per-matrix spectral-norm normalization (norm over the last two dims). + X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) + + for _ in range(steps): + A = X @ X.mT + B = b * A + c * A @ A + X = a * X + B @ X + + if transposed: + X = X.mT + + return X + + +def compute_pre_ns_update( + grad: torch.Tensor, + momentum_buffer: torch.Tensor, + momentum: float = 0.95, + nesterov: bool = True, +) -> torch.Tensor: + """ + Compute the pre-Newton-Schulz update (momentum + optional Nesterov). + + This is separated from NS so that momentum/Nesterov can be applied on shards + before all-gathering for distributed NS. + + Args: + grad: Gradient tensor. + momentum_buffer: Momentum buffer (modified in-place). + momentum: Momentum coefficient. + nesterov: Whether to use Nesterov momentum. + + Returns: + Pre-NS update tensor (same shape as grad). + """ + # SGD-style momentum: buf = momentum * buf + grad (matching Moonlight) + momentum_buffer.mul_(momentum).add_(grad) + + # Nesterov: g = g + momentum * buf, else just use buf + if nesterov: + return grad.add(momentum_buffer, alpha=momentum) + else: + return momentum_buffer.clone() + + +def compute_pre_ns_update_moe_expert( + grad: torch.Tensor, + momentum_buffer: torch.Tensor, + momentum: float = 0.95, + nesterov: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Per-expert masked momentum for stacked MoE experts. + + Identical to :func:`compute_pre_ns_update` for *active* experts (those whose + gradient slice has any nonzero element, i.e. tokens were routed to them this + step), but FREEZES the momentum of *inactive* experts -- no momentum decay + and no accumulation -- instead of decaying it. + + This distinction matters for Muon/Dion but not for Adam: Newton-Schulz + renormalizes whatever it is handed to unit spectral norm, so an inactive + expert's (decayed but nonzero) stale momentum would be blown back up into a + full-strength update in a stale direction. Freezing keeps the momentum in + reserve until the expert is routed again, matching the reference + (microsoft/dion), which sits out params/experts that received no gradient. + + This function only handles the momentum recurrence. The caller MUST also, for + inactive experts: (1) zero the Newton-Schulz output, and (2) skip the weight + update and weight decay. The returned ``active`` mask is provided for exactly + that. + + Args: + grad: Local expert gradients, shape ``[E, M, N]`` (E = local experts). + momentum_buffer: Momentum buffer ``[E, M, N]``, modified in place. + momentum: Momentum coefficient. + nesterov: Whether to use Nesterov momentum. + + Returns: + pre_ns: Pre-Newton-Schulz update, shape ``[E, M, N]``. + active: Per-expert bool mask ``[E]``; True where the expert got a gradient. + """ + # Per-expert active mask: True iff the expert's gradient slice has any nonzero + # element. An expert with no tokens routed produces an exactly-zero gradient + # slice, so this is an exact "was this expert updated" test. + active = (grad != 0).flatten(1).any(dim=1) # [E] bool + a = active.view(-1, 1, 1).to(momentum_buffer.dtype) # [E, 1, 1] in {0, 1} + + # Masked SGD-style momentum: + # active -> momentum * buf + grad (== compute_pre_ns_update) + # inactive -> buf unchanged (factor 1.0; grad is 0 there anyway) => frozen + momentum_buffer.mul_(1 - a * (1 - momentum)).add_(grad) + + # Nesterov: g = g + momentum * buf, else just use buf. Inactive experts get a + # bogus (renormalized) value here, but the caller zeros their NS output. + if nesterov: + pre_ns = grad.add(momentum_buffer, alpha=momentum) + else: + pre_ns = momentum_buffer.clone() + + return pre_ns, active diff --git a/cosmos_framework/utils/generator/dion2_with_aux_adamw.py b/cosmos_framework/utils/generator/dion2_with_aux_adamw.py new file mode 100644 index 00000000..61b56682 --- /dev/null +++ b/cosmos_framework/utils/generator/dion2_with_aux_adamw.py @@ -0,0 +1,1149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +""" +Dion2WithAuxAdamW optimizer implementation. + +DION2 (Distributed Orthogonalization) for nn.Linear weight matrices, +with auxiliary AdamW for embeddings, biases, norms, and output layers (lm_head). + +This implementation combines elements from: + +1. Microsoft DION2 (https://github.com/microsoft/dion): + - All-to-all communication pattern for efficient distributed orthogonalization + - Submatrix selection (top-k rows/columns by L1 norm) + - Error feedback for unselected parts + - Async operations for overlapping communication with computation + (TBD: not realized yet -- see the all-to-all sites in + ``_process_dion2_batch_distributed``, which currently ``wait()`` immediately) + +2. KellerJordan/Muon (https://github.com/KellerJordan/Muon): + - Newton-Schulz orthogonalization algorithm + - Quintic iteration coefficients (a=3.4445, b=-4.7750, c=2.0315) + +3. FusedAdam (cosmos_framework/utils/generator/fused_adam.py): + - DTensor handling for FSDP/TP compatibility + - Transformer Engine fused AdamW kernel + +Key differences from MuonWithAuxAdamW: +- Uses all-to-all instead of all-gather (no redundant NS computation) +- Batches params in groups of world_size for efficient distribution +- Supports submatrix selection (fraction parameter) for sparse orthogonalization +- Each rank computes NS for exactly one param per batch (truly parallel) + +Sharding -> orthogonalization (the core idea) +--------------------------------------------- +Under FSDP each weight matrix is a DTensor split across GPUs, but Newton-Schulz +(NS) is a *whole-matrix* op and needs the full matrix in one place. The two +optimizers assemble it differently: + +Weight A sharded over 4 GPUs (one row-slice each): + + GPU0:[A0] GPU1:[A1] GPU2:[A2] GPU3:[A3] + +Muon -- all-gather: every GPU rebuilds the *same* full A and runs NS on it, so +the NS work is duplicated world_size times: + + all_gather(A) -> each GPU holds full A -> every GPU runs NS(A) (N x redundant) + +DION2 -- all-to-all: process world_size matrices (A,B,C,D) together and give each +GPU one complete matrix, so the N matrices are orthogonalized in parallel with no +redundant compute: + + before (each GPU has one slice of every matrix): + + A B C D + GPU0 [A0] [B0] [C0] [D0] + GPU1 [A1] [B1] [C1] [D1] + GPU2 [A2] [B2] [C2] [D2] + GPU3 [A3] [B3] [C3] [D3] + + --- all_to_all #1 (transpose) ---> each GPU now owns one FULL matrix: + + GPU0: [A0 A1 A2 A3] = A -> NS(A) + GPU1: [B0 B1 B2 B3] = B -> NS(B) + GPU2: [C0 C1 C2 C3] = C -> NS(C) + GPU3: [D0 D1 D2 D3] = D -> NS(D) + + --- all_to_all #2 (transpose back) ---> each GPU gets its own + orthogonalized slice of every matrix, then applies the update. + +Matrices are grouped by local shard shape into batches of world_size so the +all-to-all tensors are uniform (see ``_create_dion2_batches``). With +``fraction < 1.0`` only the top-k rows/cols (by L1 norm) are sent through this +dance, and the unselected part is carried forward via error feedback. + +TODO(hsdp-replicate-redundancy): eliminate redundant orthogonalization work in +the replicate dimension under HSDP. ``world_size`` above is the *shard* mesh dim +only -- the all-to-all is confined to the FSDP shard group and the replicate +(data-parallel) dim keeps its ``Replicate`` placement. That is correct, but each +replica group independently reruns the full Newton-Schulz on identical (post +all-reduce) gradients, so the NS *compute* is duplicated ``R`` times (R = +replicate degree; e.g. 2x for the 30B-A3B run at shard_degree=64 on 128 GPUs). +This is the standard data-parallel optimizer redundancy (Adam has it too) but is +pricier here because NS is several matmuls per matrix rather than an element-wise +step. It could be removed by distributing the matrices across the *full* 2-D mesh +(shard x replicate) so every rank orthogonalizes a distinct matrix, then +broadcasting the results back across the replicate dim -- trading the duplicate +compute for an extra cross-replica collective. Worth doing only if R is large or +the optimizer step becomes a step-time bottleneck. +""" + +import math + +import torch +import torch.distributed as dist +import torch.nn as nn +import transformer_engine as te +import transformer_engine_torch as tex +from torch.distributed.tensor import DTensor, Shard + +from cosmos_framework.utils import log +from cosmos_framework.utils.misc import get_local_tensor_if_DTensor +from cosmos_framework.utils.generator.aux_optimizer_utils import ( + DEFAULT_ADAMW_MODULE_KEYWORDS, + compute_pre_ns_update, + compute_pre_ns_update_moe_expert, + split_orthogonalizable_params, + zeropower_via_newtonschulz5, + zeropower_via_newtonschulz5_batched, +) + + +class Dion2WithAuxAdamW(torch.optim.Optimizer): + """ + Dion2WithAuxAdamW optimizer. + + Uses DION2 (Distributed Orthogonalization) for nn.Linear hidden weight matrices, + and AdamW for embeddings, biases, layer norms, and output heads (lm_head). + + Key features: + - All-to-all communication for efficient distributed NS (no redundant compute) + - Submatrix selection: only orthogonalize top-k rows/columns + - Error feedback: maintains momentum for unselected parts + - Batched processing: handles world_size params per batch + + Args: + params: Iterable of parameters to optimize. + lr: Base learning rate. + muon_momentum: Momentum coefficient for Muon/DION2. + muon_lr_scale: Scale factor for Muon LR adjustment. + ns_steps: Number of Newton-Schulz iterations. + nesterov: Whether to use Nesterov momentum. + fraction: Fraction of rows/columns to orthogonalize (0 < fraction <= 1). + ef_decay: Error feedback decay factor for selected submatrix. + weight_decay: Weight decay for all parameters. + adam_betas: Beta coefficients for the auxiliary AdamW side. + eps: Epsilon for AdamW numerical stability. + use_distributed: Whether to use distributed operations. + """ + + def __init__( + self, + params, + lr: float = 1e-4, + muon_momentum: float = 0.95, + muon_lr_scale: float = 0.2, + ns_steps: int = 5, + nesterov: bool = True, + fraction: float = 1.0, # 1.0 = full matrix, <1.0 = submatrix selection + ef_decay: float = 0.95, + weight_decay: float = 0.1, + adam_betas: tuple[float, float] = (0.9, 0.95), + eps: float = 1e-8, + use_distributed: bool = True, + capturable: bool = False, + master_weights: bool = False, + adamw_module_keywords: tuple[str, ...] | None = None, + expert_param_keywords: tuple[str, ...] | None = None, + **kwargs, + ): + if kwargs: + ignored_keys = list(kwargs.keys()) + expected_ignored = {"fused", "keys_to_select", "adamw_betas", "adamw_eps"} + unexpected = set(ignored_keys) - expected_ignored + if unexpected: + import warnings + + warnings.warn(f"Dion2WithAuxAdamW ignoring unexpected kwargs: {unexpected}") + + if not (0.0 < fraction <= 1.0): + raise ValueError(f"fraction must be in (0, 1], got {fraction}") + + # Master weights requires capturable mode + if master_weights and not capturable: + raise RuntimeError("Master weights is currently only supported with capturable=True.") + + # Store hyperparameters + # Note: lr is accessed via property that reads from param_groups + # to support LR schedulers (which update param_groups[X]["lr"]) + self.wd = weight_decay + self.muon_momentum = muon_momentum + self.muon_lr_scale = muon_lr_scale + self.ns_steps = ns_steps + self.nesterov = nesterov + self.fraction = fraction + self.ef_decay = ef_decay + self.adam_betas = tuple(adam_betas) if isinstance(adam_betas, list) else adam_betas + self.eps = eps + + # Name substrings that route an nn.Linear to AdamW (output heads). Used by + # categorize_params alongside tied-weight / vocab-shape detection. + self.adamw_module_keywords = ( + tuple(adamw_module_keywords) if adamw_module_keywords else DEFAULT_ADAMW_MODULE_KEYWORDS + ) + # Name substrings that route stacked MoE expert params ([E, M, N]) to the + # DION2 side (each expert slice orthogonalized). Empty = experts stay on + # AdamW (no behavior change). + self.expert_param_keywords = tuple(expert_param_keywords) if expert_param_keywords else () + + # Distributed settings + self.use_distributed = use_distributed and dist.is_initialized() + self._world_size = 1 + self._device_rank = 0 + self._process_group = None + self._device_mesh = None + + # Master weights settings (for mixed-precision training stability) + self.capturable = capturable + self.master_weights = master_weights + + # Parameter lists + self.dion2_params: list[nn.Parameter] = [] + self.adamw_params: list[nn.Parameter] = [] + # Stacked MoE expert params ([E, M, N]); orthogonalized per expert slice. + self.stacked_dion2_params: list[nn.Parameter] = [] + self.param_to_name: dict[nn.Parameter, str] = {} + self._dion2_batches: list[list[nn.Parameter]] = [] + + # Master weight copies (populated by _create_master_weights after categorize_params) + self._dion2_masters: list[torch.Tensor] = [] + self._adamw_masters: list[torch.Tensor] = [] + + # Transformer Engine fused Adam. The zero buffer is the noop flag required + # as the second argument of TE's multi_tensor_applier; it is a fixed + # constant here (no AMP overflow handling). + self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device="cuda") + self._multi_tensor_adam = tex.multi_tensor_adam + self._multi_tensor_adam_capturable = tex.multi_tensor_adam_capturable + self._multi_tensor_adam_capturable_master = tex.multi_tensor_adam_capturable_master + + # Initialize base optimizer. betas / eps go in the defaults so each param + # group carries them (FusedAdam convention); the AdamW step reads them + # per-group, enabling per-group overrides and exact FusedAdam parity. + defaults = dict(lr=lr, weight_decay=weight_decay, betas=self.adam_betas, eps=eps) + super().__init__(params, defaults) + + # Convert LR to tensor for capturable mode + if capturable: + for idx, group in enumerate(self.param_groups): + if len(group["params"]) == 0: + continue + device = group["params"][0].device + if isinstance(group["lr"], float): + group["lr"] = torch.tensor(group["lr"], dtype=torch.float32) + self.param_groups[idx]["lr"] = group["lr"].to(device=device) + + # id(param) -> owning param_group, so the DION2 and AdamW updates can read + # the *per-group* lr / weight_decay (honors lr_multipliers and + # disable_weight_decay_for_1d_params). With a single group it degenerates + # to a single global lr/wd, matching the original (reference) behavior. + self._param_group_map: dict[int, dict] = {} + self._adamw_param_ids: set[int] = set() + # Master weights are created lazily on the first step() (FusedAdam-style), + # so that after a checkpoint resume they are rebuilt from the *restored* + # params rather than the freshly-initialized ones. + self._masters_initialized = False + self._param_to_master: dict[int, torch.Tensor] = {} + + log.info(f"Dion2WithAuxAdamW master_weights: {master_weights} capturable: {capturable}") + + def categorize_params(self, model: nn.Module) -> None: + """ + Categorize parameters into DION2 and AdamW groups; also set up distributed + configuration from the DTensor DeviceMesh. + + DION2 is used for hidden ``nn.Linear`` weights only. Embeddings, the output + head (``lm_head`` / tied / vocab-shaped projection), biases, norms, and any + non-Linear parameter go to AdamW. See + :func:`split_orthogonalizable_params` for the architecture-agnostic + embedding / output-head detection. + """ + optimizer_param_ids = {id(p) for group in self.param_groups for p in group["params"]} + + orthogonalizable, self.adamw_params, self.param_to_name = split_orthogonalizable_params( + model, + optimizer_param_ids, + adamw_module_keywords=self.adamw_module_keywords, + expert_param_keywords=self.expert_param_keywords, + ) + + # Every optimizer param lands in exactly ONE of three disjoint buckets, + # each updated by a different function in step(): + # 1. self.dion2_params -> _step_dion2 (dense 2D Linear weights) + # 2. self.stacked_dion2_params -> _step_stacked_dion2 (3D MoE experts [E, M, N]) + # 3. self.adamw_params -> _step_adamw (embeddings/head/norms/biases/1D) + # split_orthogonalizable_params separates the orthogonalizable Linear + # weights (buckets 1+2) from everything else (bucket 3); here we further + # split the orthogonalizable set into 2D (DION2 all-to-all path) vs stacked + # 3D MoE experts (orthogonalized per expert slice in _step_stacked_dion2). + self.dion2_params = [p for p in orthogonalizable if p.ndim == 2] + self.stacked_dion2_params = [p for p in orthogonalizable if p.ndim >= 3] + + # Sort by size for load balancing + self.dion2_params = sorted(self.dion2_params, key=lambda x: x.numel(), reverse=True) + + # Setup distributed from first DTensor param + self._setup_distributed_from_params() + + # Create batches of world_size params + self._create_dion2_batches() + + dion2_numel = sum(p.numel() for p in self.dion2_params) + adamw_numel = sum(p.numel() for p in self.adamw_params) + stacked_dion2_numel = sum(p.numel() for p in self.stacked_dion2_params) + + log.info( + f"Dion2WithAuxAdamW: {len(self.dion2_params)} Muon params ({dion2_numel:,} elements), " + f"{len(self.stacked_dion2_params)} stacked-expert params ({stacked_dion2_numel:,} elements), " + f"{len(self.adamw_params)} AdamW params ({adamw_numel:,} elements), " + f"world_size={self._world_size}, {len(self._dion2_batches)} batches" + ) + + # Log Muon param details + log.info("DION2 parameters (layer name -> shape):") + for p in self.dion2_params: + name = self.param_to_name.get(p, "unknown") + log.info(f" {name}: {tuple(p.shape)}") + + # Build the param -> owning group map for per-group lr / weight_decay + # lookups during the DION2 and AdamW steps. + self._param_group_map = {} + for group in self.param_groups: + for p in group["params"]: + self._param_group_map[id(p)] = group + self._adamw_param_ids = {id(p) for p in self.adamw_params} + + # NOTE: master weights are intentionally NOT created here. They are + # created lazily on the first step() (see _maybe_init_master_weights) so + # that a checkpoint resume rebuilds them from the restored params. + + def _assert_homogeneous_sharding(self) -> None: + """Verify every DION2 param shares the first param's mesh and sharding. + + The all-to-all path caches a single distributed config (device_mesh, shard + mesh dim, shard tensor dim, world_size, process group) derived from + ``dion2_params[0]`` and applies it to *every* DION2 param. That is only + correct if all params are sharded identically. Heterogeneous sharding -- a + param on a different mesh, sharded on a different dim, or replicated while + others are sharded -- would be processed with the wrong layout and silently + corrupt the update, so it is rejected up front (once, at setup). VFM's FSDP2 + shards every 2-D weight on dim 0 with one mesh, so this is a no-op today; the + guard fails loudly if TP/EP or mixed replication is ever introduced. + """ + p0 = self.dion2_params[0] + for p in self.dion2_params[1:]: + mismatch = isinstance(p, DTensor) != isinstance(p0, DTensor) or ( + isinstance(p0, DTensor) and (p.device_mesh != p0.device_mesh or p.placements != p0.placements) + ) + if mismatch: + name = self.param_to_name.get(p, "unknown") + raise NotImplementedError( + f"DION2 requires all params to share the first param's sharding, but '{name}' " + f"differs (placements={getattr(p, 'placements', None)}). " + f"Heterogeneous sharding is unsupported." + ) + + def _setup_distributed_from_params(self) -> None: + """Extract distributed config from DTensor DeviceMesh.""" + if not self.dion2_params: + return + + # The config below is derived from dion2_params[0] and reused for every + # param, so first confirm they are all sharded identically. + self._assert_homogeneous_sharding() + + first_param = self.dion2_params[0] + if isinstance(first_param, DTensor): + device_mesh = first_param.device_mesh + placements = first_param.placements + + # Find the shard dimension in the mesh + for mesh_dim_idx, placement in enumerate(placements): + if placement.is_shard(): + self._device_mesh = device_mesh + self._world_size = device_mesh.size(mesh_dim=mesh_dim_idx) + self._device_rank = device_mesh.get_local_rank(mesh_dim=mesh_dim_idx) + self._process_group = device_mesh.get_group(mesh_dim=mesh_dim_idx) + self._shard_mesh_dim = mesh_dim_idx + self._shard_tensor_dim = placement.dim + log.info( + f"DION2 distributed setup: world_size={self._world_size}, " + f"rank={self._device_rank}, shard_dim={self._shard_tensor_dim}" + ) + return + + # Fallback: not distributed or not sharded + self._world_size = 1 + self._device_rank = 0 + + def _create_dion2_batches(self) -> None: + """ + Group Muon params by GLOBAL shape, then batch within each shape group. + + The distributed step (``_process_dion2_batch_distributed``) stacks a batch of + params into a single DTensor and redistributes it, so the batches must be + identical ACROSS ranks. Group by the *global* shape (same on every rank), NOT + the local shard shape -- under uneven sharding the local shard shape differs + per rank and would make ranks build inconsistent batches. ``batch_size = + world_size`` ensures each rank orthogonalizes exactly one param per batch. + """ + self._dion2_batches = [] + batch_size = self._world_size + + # Step 1: Group params by global shape (identical on all ranks). + shape_groups: dict[tuple, list[nn.Parameter]] = {} + for p in self.dion2_params: + shape = tuple(p.shape) + if shape not in shape_groups: + shape_groups[shape] = [] + shape_groups[shape].append(p) + + # Step 2: Create batches within each shape group + for shape, params in shape_groups.items(): + for i in range(0, len(params), batch_size): + batch = params[i : i + batch_size] + self._dion2_batches.append(batch) + + # Log batch info + num_shape_groups = len(shape_groups) + num_batches = len(self._dion2_batches) + if self._dion2_batches: + # Count batches that need padding + padded_batches = sum(1 for b in self._dion2_batches if len(b) < batch_size) + log.info( + f"DION2: {len(self.dion2_params)} params grouped into {num_shape_groups} shape groups, " + f"{num_batches} batches (world_size={batch_size}, {padded_batches} need padding)" + ) + # Log shape group details + for shape, params in shape_groups.items(): + log.info(f" Shape {shape}: {len(params)} params") + + def _base_lr_for(self, p: nn.Parameter) -> float | torch.Tensor: + """Per-group base learning rate for ``p`` (honors lr_multipliers).""" + return self._param_group_map[id(p)]["lr"] + + def _wd_for(self, p: nn.Parameter) -> float: + """Per-group weight decay for ``p`` (honors disable_weight_decay_for_1d_params).""" + return self._param_group_map[id(p)]["weight_decay"] + + def _get_adjusted_lr(self, param_shape: tuple[int, ...], base_lr: float | torch.Tensor) -> float | torch.Tensor: + """Compute adjusted learning rate based on parameter matrix size and the + owning param-group's base lr.""" + A, B = param_shape[:2] + adjusted_ratio = self.muon_lr_scale * math.sqrt(max(A, B)) + return base_lr * adjusted_ratio + + def _maybe_init_master_weights(self) -> None: + """Create FP32 master weights on first use (FusedAdam-style lazy init).""" + if self.master_weights and not self._masters_initialized: + self._create_master_weights() + + def _create_master_weights(self) -> None: + """ + Create FP32 master weight copies for mixed-precision training stability. + + Creates param_groups_master (for LowPrecisionCallback compatibility) and + indexed lists for efficient lookup during DION2/AdamW steps. + """ + # Create param_groups_master mirroring param_groups (like FusedAdam) + # This enables LowPrecisionCallback to copy masters -> params periodically + self.param_groups_master = [] + for pg in self.param_groups: + param_list = pg["params"] + self.param_groups_master.append( + { + "params": [p.clone().detach().float() if self.master_weights else None for p in param_list], + } + ) + + # Build param_id -> master mapping for efficient lookup + self._param_to_master: dict[int, torch.Tensor] = {} + for group, group_master in zip(self.param_groups, self.param_groups_master): + for p, p_master in zip(group["params"], group_master["params"]): + if p_master is not None: + self._param_to_master[id(p)] = p_master + + # Create indexed lists for DION2/AdamW step() methods + self._dion2_masters = [self._param_to_master[id(p)] for p in self.dion2_params] + self._adamw_masters = [self._param_to_master[id(p)] for p in self.adamw_params] + + muon_master_numel = sum(m.numel() for m in self._dion2_masters) + adamw_master_numel = sum(m.numel() for m in self._adamw_masters) + log.info( + f"Created FP32 master weights: {len(self._dion2_masters)} DION2 ({muon_master_numel:,} elements), " + f"{len(self._adamw_masters)} AdamW ({adamw_master_numel:,} elements)" + ) + self._masters_initialized = True + + @torch.no_grad() + def step(self, closure=None): + """Perform a single optimization step.""" + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + # Lazily build FP32 master weights from the current (possibly + # checkpoint-restored) params before the first update. + self._maybe_init_master_weights() + + # Params are split into three disjoint buckets at init (see + # categorize_params) and each is updated by a different function below: + # 1. self.dion2_params (dense 2D linears) -> _step_dion2 + # 2. self.stacked_dion2_params (3D MoE experts [E,M,N]) -> _step_stacked_dion2 + # 3. self.adamw_params (embeddings/head/norms/1D) -> _step_adamw + # Every parameter belongs to exactly one bucket, so exactly one of these + # updates it. Order does not matter (buckets are disjoint). + + # 1. Dense 2D linears: orthogonalized via all-to-all distributed Newton-Schulz. + self._step_dion2() + + # 2. MoE expert weights: each [E, M, N] param orthogonalized one expert + # slice at a time (local NS; per-expert masking for inactive experts). + self._step_stacked_dion2() + + # 3. Everything else (embeddings, lm_head, norms, biases, 1D): fused AdamW. + self._step_adamw() + + return loss + + def _step_stacked_dion2(self) -> None: + """Orthogonalize stacked MoE expert params, one expert slice at a time. + + Each param has shape ``[E, M, N]`` (E experts, each an M x N matrix). Under + FSDP2 these are sharded on the expert dim (dim 0), so every rank holds whole + expert matrices -- Newton-Schulz is therefore fully local (no all-to-all), + and is batched across the local experts via ``zeropower_via_newtonschulz5_batched``. + + NOTE (sharding assumption): the "no communication" property relies on the + expert tensor being sharded on dim 0 (the expert axis). This holds for the + FSDP2 ``fully_shard`` path used by LLM/VFM here, because FSDP2 shards every + parameter on dim 0. It is NOT guaranteed in general -- e.g. tensor/expert + parallelism could shard *within* an expert matrix (dim 1/2). That case is + unsupported and is rejected by the placement check below (fails loudly + rather than silently computing a wrong update); supporting it would require + a per-expert gather. The assumption was not exhaustively audited against + every parallelization config, which is exactly why it is enforced here. + """ + for p in self.stacked_dion2_params: + if p.grad is None: + continue + + # Validate sharding: only the expert axis (tensor dim 0) may be sharded. + if isinstance(p, DTensor): + for placement in p.placements: + if placement.is_shard() and placement.dim != 0: + raise NotImplementedError( + "Stacked-expert orthogonalization requires sharding on the expert " + f"dim (0); got placement {placement} for " + f"'{self.param_to_name.get(p, 'unknown')}'." + ) + + local_grad = get_local_tensor_if_DTensor(p.grad) + local_param = get_local_tensor_if_DTensor(p) + if local_grad.ndim != 3: + raise NotImplementedError( + f"Stacked-expert orthogonalization supports 3D params [E, M, N]; " + f"got shape {tuple(local_grad.shape)} for '{self.param_to_name.get(p, 'unknown')}'." + ) + + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p).float() + + # Per-expert masked momentum + Nesterov (element-wise over [E, M, N]). + # Active experts follow the standard mu*M + G recurrence; inactive + # experts (no gradient this step) keep their momentum frozen. ``active`` + # ([E] bool) is used below to zero the update and skip weight decay for + # inactive experts -- required because Newton-Schulz would otherwise + # renormalize their stale momentum into a full-strength spurious update. + pre_ns, active = compute_pre_ns_update_moe_expert( + local_grad, + get_local_tensor_if_DTensor(state["momentum_buffer"]), + momentum=self.muon_momentum, + nesterov=self.nesterov, + ) + + # Batched Newton-Schulz over the local experts, then zero the update for + # inactive experts (their NS result is a bogus unit-norm matrix). + ortho = zeropower_via_newtonschulz5_batched(pre_ns, steps=self.ns_steps) + ortho = ortho * active.view(-1, 1, 1).to(ortho.dtype) + + # LR scaling uses the per-expert matrix shape (M, N), shared across experts. + base_lr = self._base_lr_for(p) + wd = self._wd_for(p) + adjusted_lr = self._get_adjusted_lr(tuple(p.shape[-2:]), base_lr) + + # Per-expert weight-decay factor: (1 - base_lr*wd) for active experts, + # 1.0 (no decay) for inactive ones. Combined with the zeroed update + # above, inactive experts are left completely untouched. + if self.master_weights: + master = get_local_tensor_if_DTensor(self._param_to_master[id(p)]) + a_wd = active.view(-1, 1, 1).to(master.dtype) + master.mul_(1 - a_wd * (base_lr * wd)) + master.add_(ortho.float() * (-adjusted_lr)) + local_param.copy_(master) + else: + a_wd = active.view(-1, 1, 1).to(local_param.dtype) + local_param.mul_(1 - a_wd * (base_lr * wd)) + local_param.add_(ortho.to(local_param.dtype) * (-adjusted_lr)) + + def _step_dion2(self) -> None: + """ + DION2 step with all-to-all distributed Newton-Schulz. + + For each batch of world_size params: + 1. Compute momentum + select submatrix on local shards + 2. All-to-all to redistribute shards (each rank gets full submatrix for its param) + 3. Newton-Schulz on full submatrix (each rank does different param) + 4. All-to-all to scatter results back + 5. Apply weight decay and update (to FP32 master if enabled) + """ + if not self.dion2_params: + return + + for batch in self._dion2_batches: + self._process_dion2_batch(batch) + + def _process_dion2_batch(self, batch: list[nn.Parameter]) -> None: + """Process a single batch of params using DION2 all-to-all pattern.""" + world_size = self._world_size + + # Pad batch if needed + actual_batch_size = len(batch) + if actual_batch_size < world_size: + # Pad with the last param (will be masked out) + padding = [batch[-1]] * (world_size - actual_batch_size) + batch = batch + padding + + # Check if using DTensor (FSDP) + is_dtensor = isinstance(batch[0], DTensor) + + if is_dtensor and world_size > 1: + self._process_dion2_batch_distributed(batch, actual_batch_size) + else: + self._process_dion2_batch_single(batch, actual_batch_size) + + def _process_dion2_batch_distributed(self, batch: list[nn.Parameter], actual_batch_size: int) -> None: + """Process a batch via DTensor collectives (the "each rank orthogonalizes one + whole param" transpose), correct for uneven / non-divisible shard dims. + + Rather than a hand-rolled ``all_to_all`` + ``cat``/``narrow`` (which assumed + FSDP2 padded every local shard uniformly -- it does not; ``_local_tensor`` is + unpadded and uneven), the gather/scatter is expressed through DTensor: + + 1. Momentum + Nesterov per param, kept as a DTensor so its shard metadata + (including uneven, unpadded local sizes) is preserved. + 2. ``torch.stack`` the world_size params -> a ``[W, ...]`` DTensor; the shard + tensor dim shifts to ``shard_dim + 1``. + 3. ``redistribute`` so the PARAM axis is sharded on the FSDP shard mesh dim + -> each rank owns one whole param (forward all-to-all). + 4. Newton-Schulz on that whole param (each rank a different one). + 5. ``redistribute`` back to shard the data axis (backward all-to-all), + unstack, and apply the update to each local shard. + + DTensor owns the (possibly uneven) per-rank size bookkeeping, so this is + correct regardless of divisibility. This transpose has been validated across + even/uneven/partial/bf16/shard-dim configs by the standalone multi-GPU script + ``unit_tests/dion2_uneven_shard_dtensor_check.py``. + + TODO(dion2-optionb-unittest): provide a standalone torchrun script into a + committed, CI-run multi-GPU unit test (it currently must be launched manually + via ``torchrun`` and is not part of the automated suite). + + Submatrix selection (``fraction < 1``) is not supported on this path -- it is + unused (every config runs fraction=1.0) -- and is rejected up front. + """ + # distributed path. Will need to wire it in is a + # well-scoped change: per-param DTensor select (norms via + # ``pre.abs().sum(sharded_dim).full_tensor()`` -> top-k -> ``index_select`` + # with a Replicate index), error-feedback decay on the momentum buffer, and a + # branched apply that reuses ``_apply_submatrix_update[_master]`` (index_add + # into the selected indices). Deferred for now: every config runs fraction=1.0, + # so this path is unused and not worth the added complexity yet. Single-device + # fraction<1 still works via ``_process_dion2_batch_single``. + if self.fraction != 1.0: + raise NotImplementedError( + "DION2 distributed path supports only fraction=1.0 (full-matrix " + f"orthogonalization); got fraction={self.fraction}. Submatrix selection " + "under FSDP is not implemented yet (see TODO(dion2-fsdp-fraction))." + ) + + world_size = self._world_size + mesh = self._device_mesh + shard_mesh_dim = self._shard_mesh_dim + shard_dim = self._shard_tensor_dim + stack_axis = shard_dim + 1 # torch.stack adds a leading param axis + + # Step 1: momentum + Nesterov, kept in DTensor space (metadata preserved). + # compute_pre_ns_update does not mutate the gradient, so p.grad is passed + # directly; it mutates the (sharded DTensor) momentum buffer in place. + # + # None-grad handling: a param with no gradient this step is sat out -- + # momentum frozen (compute_pre_ns_update NOT called, so no mu-decay) and no + # update applied (skipped in the apply loop via ``active``). We cannot just + # drop the slot: the stack + redistribute all-to-alls are a fixed-size + # collective every rank must enter identically, so an inactive slot instead + # contributes a zero placeholder (same DTensor sharding/dtype) to keep the + # collective shapes uniform. Newton-Schulz on zeros stays finite (norm+1e-7) + # and the result is discarded on apply. This relies on ``p.grad is None`` + # being identical across ranks -- true for dense params, where a missing grad + # is structural (an unused param is None on every rank), not data-dependent. + pre_ns_list = [] + active = [] + for i in range(actual_batch_size): + p = batch[i] + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p).float() + if p.grad is None: + active.append(False) + pre_ns_list.append(torch.zeros_like(state["momentum_buffer"]).to(torch.bfloat16)) + continue + active.append(True) + pre_ns = compute_pre_ns_update( + p.grad, state["momentum_buffer"], momentum=self.muon_momentum, nesterov=self.nesterov + ) + # Cast to bf16 BEFORE the transpose so the two all-to-alls move bf16 (not + # fp32) -- matching the previous path's communication volume. This is + # numerically identical: Newton-Schulz casts to bf16 anyway, and bf16 + # rounding commutes with the reconstruction (round-then-gather == + # gather-then-round elementwise). + pre_ns_list.append(pre_ns.to(torch.bfloat16)) + + # Pad the param axis up to world_size (mirrors batch padding); the dummy + # entries' Newton-Schulz results are discarded on apply. + padded = pre_ns_list + [pre_ns_list[-1]] * (world_size - actual_batch_size) + + # Step 2: stack -> [W, ...]; the shard tensor dim moves to stack_axis. + stacked = torch.stack(padded, dim=0) + expected = Shard(stack_axis) + if stacked.placements[shard_mesh_dim] != expected: + raise RuntimeError( + f"DION2 expected stacked placement {expected} on mesh dim {shard_mesh_dim}, " + f"got {stacked.placements} (torch.stack should shift Shard({shard_dim}) -> " + f"Shard({stack_axis}))." + ) + + # Step 3: forward all-to-all -- shard the PARAM axis on the FSDP shard mesh dim + # (keep any other mesh-dim placements, e.g. Replicate under HSDP). Each rank + # then owns one whole param. + fwd_placements = list(stacked.placements) + fwd_placements[shard_mesh_dim] = Shard(0) + per_matrix = stacked.redistribute(mesh, fwd_placements) + full_p = per_matrix.to_local()[0] # this rank's whole param + + # Step 4: Newton-Schulz on the whole param. + ortho_p = zeropower_via_newtonschulz5(full_p, steps=self.ns_steps) + + # Step 5: backward all-to-all -- re-shard the data axis, then unstack. + ortho_dt = DTensor.from_local(ortho_p.unsqueeze(0), mesh, fwd_placements, run_check=False) + back = ortho_dt.redistribute(mesh, list(stacked.placements)) + back_local = back.to_local() # [W, , ...] + + # Apply the orthogonalized update to each real param's local shard. + for i in range(actual_batch_size): + if not active[i]: + # No gradient this step: momentum was frozen above; skip weight + # decay and the update entirely (matches the reference, which + # leaves a None-grad param completely untouched). + continue + p = batch[i] + local_param = p._local_tensor + update_local = back_local[i] + base_lr = self._base_lr_for(p) + wd = self._wd_for(p) + adjusted_lr = self._get_adjusted_lr(tuple(p.shape), base_lr) + if self.master_weights: + master = get_local_tensor_if_DTensor(self._param_to_master[id(p)]) + master.mul_(1 - base_lr * wd) + master.add_(update_local.float() * (-adjusted_lr)) + local_param.copy_(master) + else: + local_param.mul_(1 - base_lr * wd) + local_param.add_(update_local.to(local_param.dtype) * (-adjusted_lr)) + + def _process_dion2_batch_single(self, batch: list[nn.Parameter], actual_batch_size: int) -> None: + """Process batch on single device (no distribution).""" + for i, p in enumerate(batch): + if i >= actual_batch_size: + continue + + # No gradient this step -> sit the param out entirely: momentum buffer + # stays frozen (not created/decayed) and no update is applied. Newton- + # Schulz renormalizes any nonzero input back to unit norm, so feeding a + # stale/zero momentum would emit a full-strength spurious update; hence + # we skip rather than treat a missing grad as grad=0. Matches the + # reference (microsoft/dion), which filters None-grad params up front. + if p.grad is None: + continue + + grad = get_local_tensor_if_DTensor(p.grad) + param = get_local_tensor_if_DTensor(p) + + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p).float() + + mom = get_local_tensor_if_DTensor(state["momentum_buffer"]) + + # Momentum + submatrix selection. Single-device fallback is not + # sharded; pass shard_dim=1 so the selection dim is rows (-2), + # matching the apply call below. (Fixes the original ``select_dim=`` + # keyword-name bug, which raised TypeError on any single-GPU / + # non-DTensor run.) + if self.fraction == 1.0: + # Muon baseline: pre-decayed momentum (M <- mu*M + G) + Nesterov, + # full-matrix NS. compute_pre_ns_update does not mutate the + # gradient, so grad can be passed directly. + pre_ns = compute_pre_ns_update( + grad, + mom, + momentum=self.muon_momentum, + nesterov=self.nesterov, + ) + submatrix, indices = self._select_submatrix(pre_ns, state, shard_dim=1) + else: + # Dion2 Algorithm 1 (fractional): pure error-feedback accumulation + # with selective decay. No whole-buffer decay and no Nesterov -- NS + # runs on the accumulated M[K], and only the selected slice is + # decayed (by ef_decay) inside _select_submatrix. This keeps the + # unselected rows as a running residual, matching the reference. + mom.add_(grad) + submatrix, indices = self._select_submatrix(mom, state, shard_dim=1) + + # Newton-Schulz + ortho = zeropower_via_newtonschulz5(submatrix, steps=self.ns_steps) + + # Get adjusted LR / wd from the owning param-group. + base_lr = self._base_lr_for(p) + wd = self._wd_for(p) + adjusted_lr = self._get_adjusted_lr(p.shape, base_lr) + + if self.master_weights: + # Update FP32 master, then write the param directly inside + # _apply_submatrix_update_master. + master = get_local_tensor_if_DTensor(self._param_to_master[id(p)]) + master.mul_(1 - base_lr * wd) + self._apply_submatrix_update_master(master, param, ortho, indices, adjusted_lr, select_dim=-2) + else: + # Apply weight decay + param.mul_(1 - base_lr * wd) + # Apply update to selected indices + self._apply_submatrix_update(param, ortho, indices, adjusted_lr, select_dim=-2) + + def _select_submatrix( + self, + tensor: torch.Tensor, + state: dict, + shard_dim: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Select submatrix based on L1 norm (DION2 style). + + Args: + tensor: Input tensor (local shard or full matrix) + state: Optimizer state dict + shard_dim: Dimension along which tensor is sharded (for FSDP) + + Returns: + submatrix: Selected rows/columns + indices: Indices of selected rows/columns + """ + if self.fraction == 1.0: + # Full matrix, no selection + if tensor.ndim == 2: + indices = torch.arange(tensor.size(0), device=tensor.device) + else: + indices = None + # Convert to BF16 for all_to_all compatibility (tensor may be FP32 from momentum) + return tensor.to(torch.bfloat16), indices + + # Determine selection dimension (opposite of shard dim for efficiency) + # If sharded along rows, select columns; if sharded along cols, select rows + if shard_dim == 0: + select_dim = -1 # Select columns + norm_dim = -2 # Compute norm over rows + else: + select_dim = -2 # Select rows + norm_dim = -1 # Compute norm over columns + + num_select = tensor.size(select_dim) + k = max(1, int(math.ceil(self.fraction * num_select))) + + # Compute L1 norm along norm_dim + slice_norms = tensor.abs().sum(dim=norm_dim) + + # All-reduce norms across ranks so all ranks select the same indices + # This is critical for FSDP where each rank has different rows/cols + # The all-reduce sums the partial norms to get global norms + if self._process_group is not None: + dist.all_reduce(slice_norms, group=self._process_group) + + # Top-k selection (now deterministic across all ranks) + _, indices = torch.topk(slice_norms, k, dim=-1, sorted=False) + + # Extract selected submatrix + if select_dim == -2: + submatrix = tensor.index_select(dim=0, index=indices) + else: + submatrix = tensor.index_select(dim=1, index=indices) + + # Apply error feedback decay to the selected rows/cols of the momentum buffer. + # Operate on the LOCAL shard (not the DTensor wrapper): ``indices`` are + # computed from the local pre-NS slice, and the selection dimension is the + # non-sharded one (opposite of shard_dim), so local and global indices + # coincide on that axis. This keeps the buffer in the same (local) coordinate + # frame as the update applied later, and avoids an unsupported in-place + # index_copy_ on a DTensor. + if "momentum_buffer" in state and self.ef_decay < 1.0: + momentum = get_local_tensor_if_DTensor(state["momentum_buffer"]) + dim = 0 if select_dim == -2 else 1 + selected = momentum.index_select(dim=dim, index=indices) + momentum.index_copy_(dim=dim, index=indices, source=selected * self.ef_decay) + + return submatrix.to(torch.bfloat16), indices + + def _apply_submatrix_update( + self, + param: torch.Tensor, + ortho: torch.Tensor, + indices: torch.Tensor | None, + lr: float | torch.Tensor, + select_dim: int, + ) -> None: + """Apply orthogonalized update to selected indices.""" + ortho = ortho.to(param.dtype) + + if indices is None or self.fraction == 1.0: + # Full matrix update. lr may be a tensor (capturable), so scale via + # multiply rather than ``alpha=`` (Tensor.add_ only accepts a Number). + param.add_(ortho * (-lr)) + else: + # Submatrix update at selected indices + scaled_ortho = -lr * ortho + if select_dim == -2 or select_dim == 0: + param.index_add_(dim=0, index=indices, source=scaled_ortho) + else: + param.index_add_(dim=1, index=indices, source=scaled_ortho) + + def _apply_submatrix_update_master( + self, + master: torch.Tensor, + param: torch.Tensor, + ortho: torch.Tensor, + indices: torch.Tensor | None, + lr: float | torch.Tensor, + select_dim: int, + ) -> None: + """Apply orthogonalized update to FP32 master, then copy to BF16 param.""" + ortho = ortho.float() + + if indices is None or self.fraction == 1.0: + # Full matrix update. lr may be a tensor (capturable), so scale via + # multiply rather than ``alpha=`` (Tensor.add_ only accepts a Number). + master.add_(ortho * (-lr)) + else: + # Submatrix update at selected indices + scaled_ortho = -lr * ortho + if select_dim == -2 or select_dim == 0: + master.index_add_(dim=0, index=indices, source=scaled_ortho) + else: + master.index_add_(dim=1, index=indices, source=scaled_ortho) + # Write the BF16 param directly from the updated FP32 master so we do not + # depend on LowPrecisionCallback (the OptimizersContainer hides + # master_weights from it). Matches FusedAdam's in-kernel param write. + param.copy_(master) + + def _step_adamw(self) -> None: + """ + AdamW step using Transformer Engine's fused kernel. + + Iterates over param groups so each group's lr / betas / eps / weight_decay + (set by the factory's lr_multipliers and disable_weight_decay_for_1d_params) + is honored, then batches by dtype within the group. The per-group step + counter lives on ``group["step"]`` (FusedAdam-style) so it is + round-tripped by the distributed-checkpoint optimizer state dict. + + When master_weights=True and capturable=True, uses + multi_tensor_adam_capturable_master which maintains FP32 master weights. + """ + if not self.adamw_params: + return + + adam_w_mode = 1 + bias_correction = 1 + + for group in self.param_groups: + # Only the AdamW-categorized params of this group are handled here; + # the DION2-categorized params were already updated in _step_dion2. + group_params = [p for p in group["params"] if id(p) in self._adamw_param_ids] + if not group_params: + continue + + device = group_params[0].device + + # Per-group step counter stored on the group (FusedAdam convention) so + # DCP round-trips it on resume. + if group.get("step", None) is not None: + if self.capturable and not isinstance(group["step"], torch.Tensor): + group["step"] = torch.tensor(group["step"], dtype=torch.int32, device=device) + group["step"] = ( + group["step"].to(device=device) if isinstance(group["step"], torch.Tensor) else group["step"] + ) + group["step"] += 1 + else: + group["step"] = torch.tensor([1], dtype=torch.int32, device=device) if self.capturable else 1 + + if self.capturable and not isinstance(group["lr"], torch.Tensor): + group["lr"] = torch.tensor(group["lr"], dtype=torch.float32, device=device) + + lr = group["lr"] + wd = group["weight_decay"] + step = group["step"] + beta1, beta2 = group["betas"] + eps = group["eps"] + + g_16, p_16, m_16, v_16 = [], [], [], [] + g_bf, p_bf, m_bf, v_bf = [], [], [], [] + g_32, p_32, m_32, v_32 = [], [], [], [] + p_16_master, p_bf_master, p_32_master = [], [], [] + + for p in group_params: + if p.grad is None: + continue + + grad = get_local_tensor_if_DTensor(p.grad) + param = get_local_tensor_if_DTensor(p) + + state = self.state[p] + if len(state) == 0: + state["exp_avg"] = torch.zeros_like(p).float() + state["exp_avg_sq"] = torch.zeros_like(p).float() + + exp_avg = get_local_tensor_if_DTensor(state["exp_avg"]) + exp_avg_sq = get_local_tensor_if_DTensor(state["exp_avg_sq"]) + master = get_local_tensor_if_DTensor(self._param_to_master[id(p)]) if self.master_weights else None + + if p.dtype == torch.float16: + g_16.append(grad) + p_16.append(param) + m_16.append(exp_avg) + v_16.append(exp_avg_sq) + if self.master_weights: + p_16_master.append(master) + elif p.dtype == torch.bfloat16: + g_bf.append(grad) + p_bf.append(param) + m_bf.append(exp_avg) + v_bf.append(exp_avg_sq) + if self.master_weights: + p_bf_master.append(master) + elif p.dtype == torch.float32: + g_32.append(grad) + p_32.append(param) + m_32.append(exp_avg) + v_32.append(exp_avg_sq) + if self.master_weights: + p_32_master.append(master) + else: + raise RuntimeError(f"Unsupported dtype {p.dtype}") + + if self.capturable: + # The capturable-master kernel requires an inverse-scale argument; + # bf16-only training has no grad scaler, so it is a constant one. + kernel_inv_scale = torch.ones((1,), device=device, dtype=torch.float32) + dtype_batches = ( + (g_16, p_16, m_16, v_16, p_16_master), + (g_bf, p_bf, m_bf, v_bf, p_bf_master), + (g_32, p_32, m_32, v_32, p_32_master), + ) + kernel = ( + self._multi_tensor_adam_capturable_master + if self.master_weights + else self._multi_tensor_adam_capturable + ) + for g_, p_, m_, v_, pm_ in dtype_batches: + if len(g_) == 0: + continue + tensor_lists = [g_, p_, m_, v_, pm_] if self.master_weights else [g_, p_, m_, v_] + te.pytorch.optimizers.multi_tensor_applier( + kernel, + self._dummy_overflow_buf, + tensor_lists, + lr, + beta1, + beta2, + eps, + step, + adam_w_mode, + bias_correction, + wd, + kernel_inv_scale, + ) + else: + dtype_batches = ( + (g_16, p_16, m_16, v_16), + (g_bf, p_bf, m_bf, v_bf), + (g_32, p_32, m_32, v_32), + ) + for g_, p_, m_, v_ in dtype_batches: + if len(g_) == 0: + continue + te.pytorch.optimizers.multi_tensor_applier( + self._multi_tensor_adam, + self._dummy_overflow_buf, + [g_, p_, m_, v_], + lr, + beta1, + beta2, + eps, + step, + adam_w_mode, + bias_correction, + wd, + ) + + def load_state_dict(self, state_dict: dict) -> None: + """Load optimizer state. + + The optimizer state (per-param momentum / exp_avg / exp_avg_sq and the + per-group ``step``) round-trips through the base ``torch.optim.Optimizer`` + state dict, so the distributed-checkpoint container can save/restore it + with FSDP2 resharding just like FusedAdam. Master weights are *not* + checkpointed; they are rebuilt from the restored params on the next + ``step()`` (see ``_maybe_init_master_weights``). + """ + super().load_state_dict(state_dict) + + # Force master weights to be rebuilt from the (now restored) params. + self._masters_initialized = False + self.param_groups_master = None + + for group in self.param_groups: + device = group["params"][0].device if group["params"] else "cuda" + if self.capturable: + if isinstance(group["lr"], torch.Tensor): + group["lr"] = group["lr"].to(device=device) + else: + group["lr"] = torch.tensor(group["lr"], dtype=torch.float32, device=device) + if group.get("step", None) is not None and not isinstance(group["step"], torch.Tensor): + group["step"] = torch.tensor(group["step"], dtype=torch.int32, device=device) + for p in group["params"]: + state = self.state[p] + if "exp_avg" in state: + state["exp_avg"] = state["exp_avg"].float() + state["exp_avg_sq"] = state["exp_avg_sq"].float() + if "momentum_buffer" in state: + state["momentum_buffer"] = state["momentum_buffer"].float() diff --git a/cosmos_framework/utils/generator/muon_with_aux_adamw.py b/cosmos_framework/utils/generator/muon_with_aux_adamw.py new file mode 100644 index 00000000..6a05e21d --- /dev/null +++ b/cosmos_framework/utils/generator/muon_with_aux_adamw.py @@ -0,0 +1,835 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +""" +MuonWithAuxAdamW optimizer implementation. + +Muon (MomentUm Orthogonalized by Newton-schulz) for nn.Linear weight matrices, +with auxiliary AdamW for embeddings, biases, norms, and output layers (lm_head). + +This implementation combines elements from: + +1. FusedAdam (cosmos_framework/utils/generator/fused_adam.py): + - DTensor handling via `get_local_tensor_if_DTensor` for FSDP/TP compatibility + - Distributed step synchronization in `load_state_dict` for checkpoint compatibility + +2. KellerJordan/Muon (https://github.com/KellerJordan/Muon): + - Newton-Schulz orthogonalization algorithm (`zeropower_via_newtonschulz5`) + - Quintic iteration coefficients (a=3.4445, b=-4.7750, c=2.0315) + - Nesterov momentum formulation + +3. MoonshotAI/Moonlight (https://arxiv.org/pdf/2502.16982): + - SGD-style momentum: `buf = momentum * buf + grad` (not EMA-style) + - Learning rate scaling by matrix size: `adjusted_lr = 0.2 * sqrt(max(A, B)) * lr` + - `@torch.compile` decoration for kernel fusion + - Parameter separation: Muon for nn.Linear weights only, AdamW for everything else + - Distributed Newton-Schulz: all-gather gradients, NS on full matrix, scatter back + (required because NS is a matrix-level operation, not element-wise) + +Sharding -> orthogonalization (the core idea) +--------------------------------------------- +Under FSDP each weight matrix is a DTensor split across GPUs, but Newton-Schulz +(NS) is a *whole-matrix* op and needs the full matrix in one place. Muon handles +each matrix one at a time with an all-gather: every GPU rebuilds the *same* full +matrix and runs NS on it. + +Weight A sharded over 4 GPUs (one row-slice each): + + GPU0:[A0] GPU1:[A1] GPU2:[A2] GPU3:[A3] + + --- all_gather(A) ---> every GPU holds the full matrix: + + GPU0:[A0 A1 A2 A3] GPU1:[A0 A1 A2 A3] GPU2:[...] GPU3:[...] + | | | | + NS(A) NS(A) NS(A) NS(A) (all identical) + | | | | + --- each GPU slices back its own rows and applies the update to its shard --- + +This is simple and correct, but the NS compute is duplicated world_size times +(every rank orthogonalizes the same matrix). The sibling ``Dion2WithAuxAdamW`` +removes that redundancy by trading shards with all-to-all so each GPU +orthogonalizes a *different* matrix in parallel (see its module docstring). +""" + +import math + +import torch +import torch.distributed as dist +import torch.nn as nn +import transformer_engine as te +import transformer_engine_torch as tex +from torch.distributed.tensor import DTensor, Replicate + +from cosmos_framework.utils import log +from cosmos_framework.utils.misc import get_local_tensor_if_DTensor +from cosmos_framework.utils.generator.aux_optimizer_utils import ( + DEFAULT_ADAMW_MODULE_KEYWORDS, + compute_pre_ns_update, + compute_pre_ns_update_moe_expert, + split_orthogonalizable_params, + zeropower_via_newtonschulz5, + zeropower_via_newtonschulz5_batched, +) + + +class MuonWithAuxAdamW(torch.optim.Optimizer): + """ + MuonWithAuxAdamW optimizer. + + Uses Muon (MomentUm Orthogonalized by Newton-schulz) for nn.Linear hidden weight matrices, + and AdamW for embeddings, biases, layer norms, and output heads (lm_head). + + See module docstring for full attribution of borrowed components. + + Args: + params: Iterable of parameters to optimize. + lr: Base learning rate. Muon scales this by muon_lr_scale*sqrt(max(A,B)), AdamW uses directly. + muon_momentum: Momentum coefficient for Muon. + muon_lr_scale: Scale factor for Muon LR adjustment. Final LR = muon_lr_scale * sqrt(max(A,B)) * lr. + ns_steps: Number of Newton-Schulz iterations. + nesterov: Whether to use Nesterov momentum for Muon. + weight_decay: Weight decay for all parameters. + adam_betas: Beta coefficients for the auxiliary AdamW side (matches VFM convention). + eps: Epsilon for AdamW numerical stability. + use_distributed: Whether to sync step counters across ranks when loading checkpoints. + """ + + def __init__( + self, + params, + lr: float = 1e-4, + muon_momentum: float = 0.95, + muon_lr_scale: float = 0.2, + ns_steps: int = 5, + nesterov: bool = True, + weight_decay: float = 0.1, + adam_betas: tuple[float, float] = (0.9, 0.95), + eps: float = 1e-8, + use_distributed: bool = True, + capturable: bool = False, + master_weights: bool = False, + adamw_module_keywords: tuple[str, ...] | None = None, + expert_param_keywords: tuple[str, ...] | None = None, + **kwargs, # Absorb VFM-specific args (fused, keys_to_select, etc.) + ): + # Log any ignored kwargs for debugging + if kwargs: + ignored_keys = list(kwargs.keys()) + # These are expected VFM args that we silently ignore + expected_ignored = {"fused", "keys_to_select", "adamw_betas", "adamw_eps"} + unexpected = set(ignored_keys) - expected_ignored + if unexpected: + import warnings + + warnings.warn(f"MuonWithAuxAdamW ignoring unexpected kwargs: {unexpected}") + + # Master weights requires capturable mode + if master_weights and not capturable: + raise RuntimeError("Master weights is currently only supported with capturable=True.") + + # Store shared hyperparameters + # Note: lr is accessed via property that reads from param_groups + # to support LR schedulers (which update param_groups[X]["lr"]) + self.wd = weight_decay + + # Store Muon-specific hyperparameters + self.muon_momentum = muon_momentum + self.muon_lr_scale = muon_lr_scale + self.ns_steps = ns_steps + self.nesterov = nesterov + + # Name substrings that route an nn.Linear to AdamW (output heads). Used by + # categorize_params alongside tied-weight / vocab-shape detection. + self.adamw_module_keywords = ( + tuple(adamw_module_keywords) if adamw_module_keywords else DEFAULT_ADAMW_MODULE_KEYWORDS + ) + # Name substrings that route stacked MoE expert params ([E, M, N]) to the + # Muon side (each expert slice orthogonalized). Empty = experts stay on + # AdamW (no behavior change). + self.expert_param_keywords = tuple(expert_param_keywords) if expert_param_keywords else () + + # Store AdamW-specific hyperparameters + self.adam_betas = tuple(adam_betas) if isinstance(adam_betas, list) else adam_betas + self.eps = eps + + # Distributed settings + self.use_distributed = use_distributed and dist.is_initialized() + + # Master weights settings (for mixed-precision training stability) + self.capturable = capturable + self.master_weights = master_weights + + # Parameter lists (populated by categorize_params) + self.muon_params: list[nn.Parameter] = [] + self.adamw_params: list[nn.Parameter] = [] + # Stacked MoE expert params ([E, M, N]); orthogonalized per expert slice. + self.stacked_muon_params: list[nn.Parameter] = [] + self.param_to_name: dict[nn.Parameter, str] = {} + + # Master weight copies (populated by _create_master_weights after categorize_params) + self._muon_masters: list[torch.Tensor] = [] + self._adamw_masters: list[torch.Tensor] = [] + + # Transformer Engine fused Adam for AdamW params. The zero buffer is the + # noop flag required as the second argument of TE's multi_tensor_applier; + # it is a fixed constant here (no AMP overflow handling). + self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device="cuda") + self._multi_tensor_adam = tex.multi_tensor_adam + self._multi_tensor_adam_capturable = tex.multi_tensor_adam_capturable + self._multi_tensor_adam_capturable_master = tex.multi_tensor_adam_capturable_master + + # Initialize base optimizer. betas / eps go in the defaults so each param + # group carries them (FusedAdam convention); the AdamW step reads them + # per-group, enabling per-group overrides and exact FusedAdam parity. + defaults = dict(lr=lr, weight_decay=weight_decay, betas=self.adam_betas, eps=eps) + super().__init__(params, defaults) + + # Convert LR to tensor for capturable mode + if capturable: + for idx, group in enumerate(self.param_groups): + if len(group["params"]) == 0: + continue + device = group["params"][0].device + if isinstance(group["lr"], float): + group["lr"] = torch.tensor(group["lr"], dtype=torch.float32) + self.param_groups[idx]["lr"] = group["lr"].to(device=device) + + # id(param) -> owning param_group, so the Muon and AdamW updates can read + # the *per-group* lr / weight_decay. This is what lets the factory's + # lr_multipliers and disable_weight_decay_for_1d_params flow through. With + # a single param group it degenerates to a single global lr/wd, matching + # the original (reference) behavior. Populated by categorize_params. + self._param_group_map: dict[int, dict] = {} + self._adamw_param_ids: set[int] = set() + # Master weights are created lazily on the first step() (FusedAdam-style), + # so that after a checkpoint resume they are rebuilt from the *restored* + # params rather than the freshly-initialized ones. + self._masters_initialized = False + self._param_to_master: dict[int, torch.Tensor] = {} + + log.info(f"MuonWithAuxAdamW master_weights: {master_weights} capturable: {capturable}") + + def categorize_params(self, model: nn.Module) -> None: + """ + Categorize parameters into Muon and AdamW groups. + + Muon is used for hidden ``nn.Linear`` weights only. Embeddings, the output + head (``lm_head`` / tied / vocab-shaped projection), biases, norms, and any + non-Linear parameter go to AdamW. See + :func:`split_orthogonalizable_params` for the architecture-agnostic + embedding / output-head detection. + + Args: + model: The model (typically the trainable ``net``) to categorize. + """ + # Only categorize params that are in this optimizer's param_groups. This + # respects keys_to_select filtering - params not passed to __init__ should + # not be categorized, otherwise state_dict() would fail with KeyError when + # mapping state entries to param indices. + optimizer_param_ids = {id(p) for group in self.param_groups for p in group["params"]} + + orthogonalizable, self.adamw_params, self.param_to_name = split_orthogonalizable_params( + model, + optimizer_param_ids, + adamw_module_keywords=self.adamw_module_keywords, + expert_param_keywords=self.expert_param_keywords, + ) + + # Split orthogonalizable params into 2-D (regular Linear weights -> Muon) and + # stacked 3-D MoE experts ([E, M, N] -> orthogonalized per expert slice). + self.muon_params = [p for p in orthogonalizable if p.ndim == 2] + self.stacked_muon_params = [p for p in orthogonalizable if p.ndim >= 3] + + # Sort Muon params by size (largest first) for distributed load balancing + self.muon_params = sorted(self.muon_params, key=lambda x: x.numel(), reverse=True) + + # Check if using DTensor (FSDP) - determines whether we use distributed NS + self._params_are_dtensor = isinstance(self.muon_params[0], DTensor) if self.muon_params else False + + muon_numel = sum(p.numel() for p in self.muon_params) + adamw_numel = sum(p.numel() for p in self.adamw_params) + stacked_muon_numel = sum(p.numel() for p in self.stacked_muon_params) + + log.info( + f"MuonWithAuxAdamW: {len(self.muon_params)} Muon params ({muon_numel:,} elements), " + f"{len(self.stacked_muon_params)} stacked-expert params ({stacked_muon_numel:,} elements), " + f"{len(self.adamw_params)} AdamW params ({adamw_numel:,} elements)" + f"{', using distributed NS (DTensor/FSDP detected)' if self._params_are_dtensor else ''}" + ) + + # Log Muon param details (layer names and shapes) + log.info("Muon parameters (layer name -> shape):") + for p in self.muon_params: + name = self.param_to_name.get(p, "unknown") + log.info(f" {name}: {tuple(p.shape)}") + + # Build the param -> owning group map for per-group lr / weight_decay + # lookups during the Muon and AdamW steps. + self._param_group_map = {} + for group in self.param_groups: + for p in group["params"]: + self._param_group_map[id(p)] = group + self._adamw_param_ids = {id(p) for p in self.adamw_params} + + # NOTE: master weights are intentionally NOT created here. They are + # created lazily on the first step() (see _maybe_init_master_weights) so + # that a checkpoint resume rebuilds them from the restored params. + + def _base_lr_for(self, p: nn.Parameter) -> float | torch.Tensor: + """Per-group base learning rate for ``p`` (honors lr_multipliers).""" + return self._param_group_map[id(p)]["lr"] + + def _wd_for(self, p: nn.Parameter) -> float: + """Per-group weight decay for ``p`` (honors disable_weight_decay_for_1d_params).""" + return self._param_group_map[id(p)]["weight_decay"] + + def _get_adjusted_lr(self, param_shape: tuple[int, ...], base_lr: float | torch.Tensor) -> float | torch.Tensor: + """ + Compute adjusted learning rate based on parameter matrix size. + + Based on Moonlight: adjusted_lr = muon_lr_scale * sqrt(max(A, B)) * base_lr + + Args: + param_shape: Shape of the parameter tensor. + base_lr: The owning param-group's learning rate (after any + lr_multiplier). Muon's matrix-size scaling layers on top of it. + + Returns: + Adjusted learning rate for this parameter. + """ + A, B = param_shape[:2] + adjusted_ratio = self.muon_lr_scale * math.sqrt(max(A, B)) + return base_lr * adjusted_ratio + + def _maybe_init_master_weights(self) -> None: + """Create FP32 master weights on first use (FusedAdam-style lazy init).""" + if self.master_weights and not self._masters_initialized: + self._create_master_weights() + + def _create_master_weights(self) -> None: + """ + Create FP32 master weight copies for mixed-precision training stability. + + Creates param_groups_master (for LowPrecisionCallback compatibility) and + indexed lists for efficient lookup during Muon/AdamW steps. + """ + # Create param_groups_master mirroring param_groups (like FusedAdam) + # This enables LowPrecisionCallback to copy masters -> params periodically + self.param_groups_master = [] + for pg in self.param_groups: + param_list = pg["params"] + self.param_groups_master.append( + { + "params": [p.clone().detach().float() if self.master_weights else None for p in param_list], + } + ) + + # Build param_id -> master mapping for efficient lookup + self._param_to_master: dict[int, torch.Tensor] = {} + for group, group_master in zip(self.param_groups, self.param_groups_master): + for p, p_master in zip(group["params"], group_master["params"]): + if p_master is not None: + self._param_to_master[id(p)] = p_master + + # Create indexed lists for Muon/AdamW step() methods + self._muon_masters = [self._param_to_master[id(p)] for p in self.muon_params] + self._adamw_masters = [self._param_to_master[id(p)] for p in self.adamw_params] + + muon_master_numel = sum(m.numel() for m in self._muon_masters) + adamw_master_numel = sum(m.numel() for m in self._adamw_masters) + log.info( + f"Created FP32 master weights: {len(self._muon_masters)} Muon ({muon_master_numel:,} elements), " + f"{len(self._adamw_masters)} AdamW ({adamw_master_numel:,} elements)" + ) + self._masters_initialized = True + + @torch.no_grad() + def step(self, closure=None): + """ + Perform a single optimization step. + + Args: + closure: A closure that reevaluates the model and returns the loss. + + Returns: + Loss value if closure is provided, else None. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + # Lazily build FP32 master weights from the current (possibly + # checkpoint-restored) params before the first update. + self._maybe_init_master_weights() + + # Muon updates (with distributed NS for FSDP) + self._step_muon() + + # Stacked MoE expert updates (orthogonalize each expert slice) + self._step_stacked_muon() + + # AdamW updates + self._step_adamw() + + return loss + + def _step_stacked_muon(self) -> None: + """Orthogonalize stacked MoE expert params, one expert slice at a time. + + Each param has shape ``[E, M, N]`` (E experts, each an M x N matrix). Under + FSDP2 these are sharded on the expert dim (dim 0), so every rank holds whole + expert matrices -- Newton-Schulz is therefore fully local (no all-gather), + and is batched across the local experts via ``zeropower_via_newtonschulz5_batched``. + + NOTE (sharding assumption): the "no communication" property relies on the + expert tensor being sharded on dim 0 (the expert axis). This holds for the + FSDP2 ``fully_shard`` path used by LLM/VFM here, because FSDP2 shards every + parameter on dim 0. It is NOT guaranteed in general -- e.g. tensor/expert + parallelism could shard *within* an expert matrix (dim 1/2). That case is + unsupported and is rejected by the placement check below (fails loudly + rather than silently computing a wrong update); supporting it would require + a per-expert gather. The assumption was not exhaustively audited against + every parallelization config, which is exactly why it is enforced here. + """ + for p in self.stacked_muon_params: + if p.grad is None: + continue + + # Validate sharding: only the expert axis (tensor dim 0) may be sharded. + if isinstance(p, DTensor): + for placement in p.placements: + if placement.is_shard() and placement.dim != 0: + raise NotImplementedError( + "Stacked-expert orthogonalization requires sharding on the expert " + f"dim (0); got placement {placement} for " + f"'{self.param_to_name.get(p, 'unknown')}'." + ) + + local_grad = get_local_tensor_if_DTensor(p.grad) + local_param = get_local_tensor_if_DTensor(p) + if local_grad.ndim != 3: + raise NotImplementedError( + f"Stacked-expert orthogonalization supports 3D params [E, M, N]; " + f"got shape {tuple(local_grad.shape)} for '{self.param_to_name.get(p, 'unknown')}'." + ) + + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p).float() + + # Per-expert masked momentum + Nesterov (element-wise over [E, M, N]). + # Active experts follow the standard mu*M + G recurrence; inactive + # experts (no gradient this step) keep their momentum frozen. ``active`` + # ([E] bool) is used below to zero the update and skip weight decay for + # inactive experts -- required because Newton-Schulz would otherwise + # renormalize their stale momentum into a full-strength spurious update. + pre_ns, active = compute_pre_ns_update_moe_expert( + local_grad, + get_local_tensor_if_DTensor(state["momentum_buffer"]), + momentum=self.muon_momentum, + nesterov=self.nesterov, + ) + + # Batched Newton-Schulz over the local experts, then zero the update for + # inactive experts (their NS result is a bogus unit-norm matrix). + ortho = zeropower_via_newtonschulz5_batched(pre_ns, steps=self.ns_steps) + ortho = ortho * active.view(-1, 1, 1).to(ortho.dtype) + + # LR scaling uses the per-expert matrix shape (M, N), shared across experts. + base_lr = self._base_lr_for(p) + wd = self._wd_for(p) + adjusted_lr = self._get_adjusted_lr(tuple(p.shape[-2:]), base_lr) + + # Per-expert weight-decay factor: (1 - base_lr*wd) for active experts, + # 1.0 (no decay) for inactive ones. Combined with the zeroed update + # above, inactive experts are left completely untouched. + if self.master_weights: + master = get_local_tensor_if_DTensor(self._param_to_master[id(p)]) + a_wd = active.view(-1, 1, 1).to(master.dtype) + master.mul_(1 - a_wd * (base_lr * wd)) + master.add_(ortho.float() * (-adjusted_lr)) + local_param.copy_(master) + else: + a_wd = active.view(-1, 1, 1).to(local_param.dtype) + local_param.mul_(1 - a_wd * (base_lr * wd)) + local_param.add_(ortho.to(local_param.dtype) * (-adjusted_lr)) + + def _step_muon(self) -> None: + """ + Muon step with distributed Newton-Schulz (following Moonlight paper). + + For DTensor/FSDP parameters: + 1. Apply momentum + Nesterov on local shards (element-wise, shard-safe) + 2. All-gather shards to reconstruct full pre-NS gradient + 3. Run Newton-Schulz on full matrix (must be done globally) + 4. Scatter back to get local orthogonalized update + 5. Apply weight decay and update to local shard (or FP32 master if enabled) + + For non-DTensor parameters, runs single-device Muon. + + When master_weights=True, updates are applied to FP32 masters and + then copied back to BF16 params for numerical stability. + """ + for idx, p in enumerate(self.muon_params): + if p.grad is None: + continue + + # Sharding info drives the routing below: a non-empty list means the + # param is FSDP-sharded (-> distributed NS); an empty list (a plain + # tensor, or a fully-replicated DTensor) takes the single-device path. + shard_info: list[tuple[int, int]] = [] + + if isinstance(p, DTensor): + # Local shard of the parameter (used to write the update in Step 5). + # The gradient / momentum stay as DTensors and are handled below. + local_param = p._local_tensor + + # Get sharding info from DTensor (FSDP typically shards dim 0 / rows). + # Each entry: (mesh_dim, tensor_shard_dim) + device_mesh = p.device_mesh + placements = p.placements + for mesh_dim_idx, placement in enumerate(placements): + if placement.is_shard(): + shard_info.append((mesh_dim_idx, placement.dim)) + + if not shard_info: + # Fully replicated DTensor (e.g. a DDP-like config where + # placements == (Replicate(),)). Legitimate -- fall back to + # single-device NS and log at info level (a warning here would + # needlessly alarm anyone reading the logs). + param_name = self.param_to_name.get(p, "unknown") + log.info( + f"DTensor '{param_name}' has no Shard placement (placements={placements}); " + "running single-device Newton-Schulz (replicated param)." + ) + + if shard_info: + # Distributed Muon: all-gather → NS → scatter + # Handles both 1D sharding (FSDP) and 2D sharding (FSDP + TP) + + # Initialize state with local shard shape + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p).float() + + # Step 1: Momentum + Nesterov, kept in DTensor space so the buffer + # stays sharded (memory-efficient) and the gather below is + # padding-aware. These ops are element-wise, so DTensor runs them on + # the local shards under the hood. compute_pre_ns_update does not + # mutate the gradient, so p.grad can be passed directly. + pre_ns_dtensor = compute_pre_ns_update( + p.grad, + state["momentum_buffer"], + momentum=self.muon_momentum, + nesterov=self.nesterov, + ) + + # Step 2: Reconstruct the full matrix from its shards. + # + # ``full_tensor()`` is DTensor's own all-gather (a Shard -> Replicate + # redistribute). Because the DTensor carries the logical shape, it + # rebuilds exactly ``p.shape`` and strips any FSDP2 padding of a shard + # dim that is not divisible by the mesh size. A hand-rolled + # ``all_gather`` + ``cat`` would keep those padding rows and feed a + # wrong-sized / misaligned matrix into Newton-Schulz (the uneven-shard + # bug this path used to have). + # + # PERF TODO: this is a *synchronous* all-gather (we block until the + # full matrix is reconstructed before running Newton-Schulz). If the + # optimizer step ever becomes a bottleneck, explore pipelining: overlap + # the gather for parameter (i) with the Newton-Schulz compute for + # parameter (i-1) across the muon_params loop. + # + # Cast to bf16 BEFORE the all-gather so the collective moves bf16 (not + # fp32), halving its communication volume (matches the dion2 path). This + # is numerically identical: Newton-Schulz casts to bf16 anyway, and bf16 + # rounding commutes with the reconstruction (round-then-gather == + # gather-then-round elementwise). Keep in sync with the input dtype of + # ``zeropower_via_newtonschulz5``. The fp32 momentum buffer is untouched + # (only this transient gathered copy is downcast). + full_pre_ns = pre_ns_dtensor.bfloat16().full_tensor() + + # Sanity check: full_tensor() must return the logical (unpadded) shape. + expected_shape = tuple(p.shape) + actual_shape = tuple(full_pre_ns.shape) + assert actual_shape == expected_shape, ( + f"Failed to reconstruct full matrix for '{self.param_to_name.get(p, 'unknown')}': " + f"got {actual_shape}, expected {expected_shape}" + ) + + # Log shapes on first step to confirm distributed NS is working + if "logged_shapes" not in state: + state["logged_shapes"] = True + log.info( + f"Muon distributed NS: '{self.param_to_name.get(p, 'unknown')}' " + f"local={tuple(local_param.shape)} → full={actual_shape} (verified)" + ) + + # Step 3: Newton-Schulz on full matrix + full_ortho = zeropower_via_newtonschulz5(full_pre_ns, steps=self.ns_steps) + + # Step 4: Scatter the orthogonalized full matrix back to p's sharding. + # + # full_ortho is identical on every rank, so we treat it as Replicate + # and redistribute to p's placements. Replicate -> Shard is pure local + # slicing (no communication) and re-applies the exact same (possibly + # padded) shard layout p has, so local_update lines up with the local + # shard / FP32 master. + ortho_dtensor = DTensor.from_local( + full_ortho, + device_mesh, + [Replicate()] * device_mesh.ndim, + run_check=False, + ) + local_update = ortho_dtensor.redistribute(device_mesh, placements).to_local() + + # Step 5: Apply weight decay and update + # Use full shape for LR scaling (not local shard shape) and the + # owning param-group's lr/wd (honors lr_multipliers / WD grouping). + base_lr = self._base_lr_for(p) + wd = self._wd_for(p) + adjusted_lr = self._get_adjusted_lr(full_pre_ns.shape, base_lr) + + if self.master_weights: + # Update FP32 master, then write the BF16 param directly so we + # do not depend on LowPrecisionCallback (the OptimizersContainer + # hides master_weights from it). Matches FusedAdam, whose kernel + # writes the param in-place. + # + # NOTE: adjusted_lr may be a tensor (capturable mode), so we + # scale via multiply rather than ``alpha=`` (Tensor.add_ only + # accepts a Python-number alpha). + master = get_local_tensor_if_DTensor(self._muon_masters[idx]) + master.mul_(1 - base_lr * wd) + master.add_(local_update.float() * (-adjusted_lr)) + local_param.copy_(master) + else: + local_param.mul_(1 - base_lr * wd) + local_param.add_(local_update * (-adjusted_lr)) + + else: + # Single-device Muon (non-FSDP or non-sharded) + grad = get_local_tensor_if_DTensor(p.grad) + param = get_local_tensor_if_DTensor(p) + + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p).float() + + # Compute pre-NS update (momentum + Nesterov). compute_pre_ns_update + # does not mutate the gradient, so grad can be passed directly. + pre_ns = compute_pre_ns_update( + grad, + get_local_tensor_if_DTensor(state["momentum_buffer"]), + momentum=self.muon_momentum, + nesterov=self.nesterov, + ) + + # Newton-Schulz orthogonalization + update = zeropower_via_newtonschulz5(pre_ns, steps=self.ns_steps) + + # Get adjusted LR based on matrix size (Moonlight scaling), using + # the owning param-group's lr/wd. + base_lr = self._base_lr_for(p) + wd = self._wd_for(p) + adjusted_lr = self._get_adjusted_lr(p.shape, base_lr) + + # Apply weight decay (using base LR) and update (using adjusted LR) + if self.master_weights: + # Update FP32 master, then write the BF16 param directly (see + # the distributed branch above for rationale). adjusted_lr may be + # a tensor (capturable), so scale via multiply, not ``alpha=``. + master = get_local_tensor_if_DTensor(self._muon_masters[idx]) + master.mul_(1 - base_lr * wd) + master.add_(update.float() * (-adjusted_lr)) + param.copy_(master) + else: + param.mul_(1 - base_lr * wd) + param.add_(update * (-adjusted_lr)) + + def _step_adamw(self) -> None: + """ + AdamW step using Transformer Engine's fused multi-tensor Adam kernel. + + Iterates over param groups so each group's lr / betas / eps / weight_decay + (set by the factory's lr_multipliers and disable_weight_decay_for_1d_params) + is honored, then batches by dtype (fp16, bf16, fp32) within the group. The + per-group step counter lives on ``group["step"]`` (FusedAdam-style) so + it is round-tripped by the distributed-checkpoint optimizer state dict. + + When master_weights=True and capturable=True, uses + multi_tensor_adam_capturable_master which maintains FP32 master weights. + """ + if not self.adamw_params: + return + + adam_w_mode = 1 # Decoupled weight decay + bias_correction = 1 + + for group in self.param_groups: + # Only the AdamW-categorized params of this group are handled here; + # the Muon-categorized params were already updated in _step_muon. + group_params = [p for p in group["params"] if id(p) in self._adamw_param_ids] + if not group_params: + continue + + device = group_params[0].device + + # Per-group step counter stored on the group (FusedAdam convention) so + # DCP round-trips it on resume. + if group.get("step", None) is not None: + if self.capturable and not isinstance(group["step"], torch.Tensor): + group["step"] = torch.tensor(group["step"], dtype=torch.int32, device=device) + group["step"] = ( + group["step"].to(device=device) if isinstance(group["step"], torch.Tensor) else group["step"] + ) + group["step"] += 1 + else: + group["step"] = torch.tensor([1], dtype=torch.int32, device=device) if self.capturable else 1 + + if self.capturable and not isinstance(group["lr"], torch.Tensor): + group["lr"] = torch.tensor(group["lr"], dtype=torch.float32, device=device) + + lr = group["lr"] + wd = group["weight_decay"] + step = group["step"] + beta1, beta2 = group["betas"] + eps = group["eps"] + + # Batch parameters by dtype for multi-tensor apply. + g_16, p_16, m_16, v_16 = [], [], [], [] + g_bf, p_bf, m_bf, v_bf = [], [], [], [] + g_32, p_32, m_32, v_32 = [], [], [], [] + p_16_master, p_bf_master, p_32_master = [], [], [] + + for p in group_params: + if p.grad is None: + continue + + grad = get_local_tensor_if_DTensor(p.grad) + param = get_local_tensor_if_DTensor(p) + + state = self.state[p] + if len(state) == 0: + state["exp_avg"] = torch.zeros_like(p).float() + state["exp_avg_sq"] = torch.zeros_like(p).float() + + exp_avg = get_local_tensor_if_DTensor(state["exp_avg"]) + exp_avg_sq = get_local_tensor_if_DTensor(state["exp_avg_sq"]) + master = get_local_tensor_if_DTensor(self._param_to_master[id(p)]) if self.master_weights else None + + if p.dtype == torch.float16: + g_16.append(grad) + p_16.append(param) + m_16.append(exp_avg) + v_16.append(exp_avg_sq) + if self.master_weights: + p_16_master.append(master) + elif p.dtype == torch.bfloat16: + g_bf.append(grad) + p_bf.append(param) + m_bf.append(exp_avg) + v_bf.append(exp_avg_sq) + if self.master_weights: + p_bf_master.append(master) + elif p.dtype == torch.float32: + g_32.append(grad) + p_32.append(param) + m_32.append(exp_avg) + v_32.append(exp_avg_sq) + if self.master_weights: + p_32_master.append(master) + else: + raise RuntimeError(f"Unsupported dtype {p.dtype} for fused AdamW") + + if self.capturable: + # The capturable-master kernel requires an inverse-scale argument; + # bf16-only training has no grad scaler, so it is a constant one. + kernel_inv_scale = torch.ones((1,), device=device, dtype=torch.float32) + dtype_batches = ( + (g_16, p_16, m_16, v_16, p_16_master), + (g_bf, p_bf, m_bf, v_bf, p_bf_master), + (g_32, p_32, m_32, v_32, p_32_master), + ) + kernel = ( + self._multi_tensor_adam_capturable_master + if self.master_weights + else self._multi_tensor_adam_capturable + ) + for g_, p_, m_, v_, pm_ in dtype_batches: + if len(g_) == 0: + continue + tensor_lists = [g_, p_, m_, v_, pm_] if self.master_weights else [g_, p_, m_, v_] + te.pytorch.optimizers.multi_tensor_applier( + kernel, + self._dummy_overflow_buf, + tensor_lists, + lr, + beta1, + beta2, + eps, + step, + adam_w_mode, + bias_correction, + wd, + kernel_inv_scale, + ) + else: + dtype_batches = ( + (g_16, p_16, m_16, v_16), + (g_bf, p_bf, m_bf, v_bf), + (g_32, p_32, m_32, v_32), + ) + for g_, p_, m_, v_ in dtype_batches: + if len(g_) == 0: + continue + te.pytorch.optimizers.multi_tensor_applier( + self._multi_tensor_adam, + self._dummy_overflow_buf, + [g_, p_, m_, v_], + lr, + beta1, + beta2, + eps, + step, + adam_w_mode, + bias_correction, + wd, + ) + + def load_state_dict(self, state_dict: dict) -> None: + """Load optimizer state. + + The optimizer state (per-param momentum / exp_avg / exp_avg_sq and the + per-group ``step``) round-trips through the base ``torch.optim.Optimizer`` + state dict, so the distributed-checkpoint container can save/restore it + with FSDP2 resharding just like FusedAdam. Master weights are *not* + checkpointed; they are rebuilt from the restored params on the next + ``step()`` (see ``_maybe_init_master_weights``). + + This direct-call path (used outside the OptimizersContainer / DCP flow) + keeps the moments in FP32 and normalizes the capturable LR/step tensors. + """ + super().load_state_dict(state_dict) + + # Force master weights to be rebuilt from the (now restored) params. + self._masters_initialized = False + self.param_groups_master = None + + for group in self.param_groups: + device = group["params"][0].device if group["params"] else "cuda" + if self.capturable: + if isinstance(group["lr"], torch.Tensor): + group["lr"] = group["lr"].to(device=device) + else: + group["lr"] = torch.tensor(group["lr"], dtype=torch.float32, device=device) + if group.get("step", None) is not None and not isinstance(group["step"], torch.Tensor): + group["step"] = torch.tensor(group["step"], dtype=torch.int32, device=device) + for p in group["params"]: + state = self.state[p] + if "exp_avg" in state: + state["exp_avg"] = state["exp_avg"].float() + state["exp_avg_sq"] = state["exp_avg_sq"].float() + if "momentum_buffer" in state: + state["momentum_buffer"] = state["momentum_buffer"].float() diff --git a/cosmos_framework/utils/generator/optimizer.py b/cosmos_framework/utils/generator/optimizer.py index 2c1c798b..fe446f37 100644 --- a/cosmos_framework/utils/generator/optimizer.py +++ b/cosmos_framework/utils/generator/optimizer.py @@ -15,6 +15,12 @@ from cosmos_framework.utils.functional.lr_scheduler import LambdaLinearScheduler, LambdaWarmUpCosineScheduler, WSDScheduler from cosmos_framework.utils import log +# Hybrid orthogonalizing optimizers (Muon / Dion2) that own their parameter +# categorization and run their own collective communication. They must be built +# as a single optimizer instance over all selected params (no per-device-mesh +# split) and need ``categorize_params`` called after construction. +_AUX_ADAMW_OPTIMIZERS = ("muonwithauxadamw", "dion2withauxadamw") + class ParamMetadata(NamedTuple): lr: float @@ -68,6 +74,23 @@ def _optimizer_cls( optimizer_kwargs["capturable"] = True optimizer_kwargs["master_weights"] = True optimizer = FusedAdam(params, **optimizer_kwargs) + elif optimizer_type.lower() == "muonwithauxadamw": + from cosmos_framework.utils.generator.muon_with_aux_adamw import MuonWithAuxAdamW + + # Muon's AdamW side is the TE-fused kernel; it is fused by construction and + # absorbs ``fused`` via **kwargs, but we pop it here to be explicit. We force + # capturable + master_weights to match FusedAdam's mixed-precision setup. + optimizer_kwargs.pop("fused", None) + optimizer_kwargs["capturable"] = True + optimizer_kwargs["master_weights"] = True + optimizer = MuonWithAuxAdamW(params, **optimizer_kwargs) + elif optimizer_type.lower() == "dion2withauxadamw": + from cosmos_framework.utils.generator.dion2_with_aux_adamw import Dion2WithAuxAdamW + + optimizer_kwargs.pop("fused", None) + optimizer_kwargs["capturable"] = True + optimizer_kwargs["master_weights"] = True + optimizer = Dion2WithAuxAdamW(params, **optimizer_kwargs) else: raise NotImplementedError(f"Optimizer {optimizer_type} not found.") return optimizer @@ -292,27 +315,49 @@ def __init__( disable_weight_decay_for_1d_params=disable_weight_decay_for_1d_params, ) - # Sub-group by device mesh so fused optimizers operate on same-mesh params. - mesh_groups: dict[str, list[tuple[nn.Parameter, ParamMetadata]]] = collections.defaultdict(list) - for param, metadata in params_with_metadata: - if hasattr(param, "device_mesh"): - # ``mesh_dim_names`` is ``tuple[str, ...] | None`` on DeviceMesh — - # fall back to ``default`` when names weren't assigned. - names = param.device_mesh.mesh_dim_names - mesh_key = "-".join(names) if names else "default" - else: - mesh_key = "default" - mesh_groups[mesh_key].append((param, metadata)) - - # Create one optimizer per mesh, each with per-LR,weight-decay param groups. - for mesh_key, mesh_params in mesh_groups.items(): - log.info(f"Building optimizer for mesh '{mesh_key}'") + if optimizer_type.lower() in _AUX_ADAMW_OPTIMIZERS: + # Muon / Dion2 categorize parameters into an orthogonalized-matrix group + # and an auxiliary AdamW group internally, and run their own collective + # communication across the FULL device mesh. They must therefore be a + # single optimizer instance over all selected params (no per-mesh split). + # + # The factory's param grouping is preserved: ``_build_optimizer_internal`` + # still produces the same (lr, weight_decay) param groups (from + # ``lr_multipliers`` / ``disable_weight_decay_for_1d_params``), and + # Muon/Dion2 read lr/weight_decay per group — degenerating to a single + # global lr/wd when there is only one group (the reference behavior). optimizer = _build_optimizer_internal( - mesh_params, + params_with_metadata, optimizer_type, **optimizer_kwargs, ) + # Categorize against the trainable network only, mirroring + # ``_build_params_with_metadata`` which optimizes the ``net`` subtree + # and skips ``net_ema``. + optimizer.categorize_params(model.net) self.optimizers.append(optimizer) + else: + # Sub-group by device mesh so fused optimizers operate on same-mesh params. + mesh_groups: dict[str, list[tuple[nn.Parameter, ParamMetadata]]] = collections.defaultdict(list) + for param, metadata in params_with_metadata: + if hasattr(param, "device_mesh"): + # ``mesh_dim_names`` is ``tuple[str, ...] | None`` on DeviceMesh — + # fall back to ``default`` when names weren't assigned. + names = param.device_mesh.mesh_dim_names + mesh_key = "-".join(names) if names else "default" + else: + mesh_key = "default" + mesh_groups[mesh_key].append((param, metadata)) + + # Create one optimizer per mesh, each with per-LR,weight-decay param groups. + for mesh_key, mesh_params in mesh_groups.items(): + log.info(f"Building optimizer for mesh '{mesh_key}'") + optimizer = _build_optimizer_internal( + mesh_params, + optimizer_type, + **optimizer_kwargs, + ) + self.optimizers.append(optimizer) log.info(f"Created {len(self.optimizers)} optimizers") @@ -323,6 +368,14 @@ def __len__(self) -> int: return len(self.optimizers) def step(self) -> None: + # NOTE: This container does not forward a ``torch.amp.GradScaler`` to the + # inner optimizers, and training here is bf16 with the scaler disabled + # (``grad_scaler_args={"enabled": False}``), so step() always runs unscaled. + # MuonWithAuxAdamW / Dion2WithAuxAdamW are intentionally bf16-only (their + # AMP grad-scaler support was removed to keep the Newton-Schulz path simple); + # only FusedAdam still carries a scaler hook. Enabling fp16 + dynamic loss + # scaling end-to-end would require re-adding scaler passthrough here and in + # the Muon/Dion2 step paths, which is deliberately not supported. for optimizer in self.optimizers: optimizer.step() diff --git a/cosmos_framework/utils/generator/torchcodec_video.py b/cosmos_framework/utils/generator/torchcodec_video.py index 6a27ac0d..3f8163d8 100644 --- a/cosmos_framework/utils/generator/torchcodec_video.py +++ b/cosmos_framework/utils/generator/torchcodec_video.py @@ -46,11 +46,16 @@ def _build_decoder( num_threads: int = 1, seek_mode: str = "exact", device: str = "cpu", + custom_frame_mappings: bytes | None = None, ) -> Any: normalized_source = _normalize_source(source) # Preserve FFmpeg/TorchCodec's 0 sentinel so callers can request automatic thread selection. num_ffmpeg_threads = 0 if num_threads == 0 else max(num_threads, 1) - kwargs: dict[str, Any] = {"seek_mode": seek_mode, "num_ffmpeg_threads": num_ffmpeg_threads} + kwargs: dict[str, Any] = {"num_ffmpeg_threads": num_ffmpeg_threads} + if custom_frame_mappings is None: + kwargs["seek_mode"] = seek_mode + else: + kwargs["custom_frame_mappings"] = custom_frame_mappings if device != "cpu": kwargs["device"] = device video_decoder_cls = _get_video_decoder_cls() @@ -107,8 +112,15 @@ def __init__( seek_mode: str = "exact", device: str = "cpu", include_dimensions: bool = False, + custom_frame_mappings: bytes | None = None, ) -> None: - self._decoder = _build_decoder(source, num_threads=num_threads, seek_mode=seek_mode, device=device) + self._decoder = _build_decoder( + source, + num_threads=num_threads, + seek_mode=seek_mode, + device=device, + custom_frame_mappings=custom_frame_mappings, + ) self.metadata = _metadata_from_frame(self._decoder, include_dimensions=include_dimensions) def __len__(self) -> int: