From 60c1bb6bed56e99a6cc3f12ff2ef419e42476699 Mon Sep 17 00:00:00 2001 From: Liang Feng Date: Sat, 24 Oct 2026 15:54:14 +0000 Subject: [PATCH 1/2] Release 2026-10-24 from i4 b3866eb30721bf68dce73380563235d39fa5a622 _source_commit: b3866eb30721bf68dce73380563235d39fa5a622-dirty _dest_commit: bf492fe2fda2108aae9b9561dc9fb9eec3fe2027 --- .file_mapping.json | 40 +- .../auxiliary/guardrail/common/presets.py | 7 +- .../guardrail/qwen3guard/qwen3guard.py | 6 + .../callbacks/data_stats_writer.py | 13 +- .../callbacks/dataloader_state.py | 13 +- .../callbacks/every_n_draw_sample.py | 20 +- .../callbacks/sampled_media_recorder.py | 2 +- .../callbacks/sampled_media_recorder_test.py | 4 +- cosmos_framework/callbacks/tokens_per_sec.py | 9 +- .../callbacks/wall_clock_checkpoint.py | 6 +- cosmos_framework/checkpoint/base.py | 4 +- cosmos_framework/checkpoint/dcp.py | 2 +- cosmos_framework/checkpoint/dcp_distill.py | 6 +- .../configs/base/defaults/model_config.py | 2 + .../base/defaults/multiview_attention.py | 15 + .../configs/base/defaults/reasoner.py | 9 +- .../base/reasoner/defaults/augmentors.py | 17 +- .../base/reasoner/defaults/callbacks.py | 11 +- .../configs/base/reasoner/defaults/config.py | 77 +- .../base/reasoner/defaults/policy_config.py | 15 +- cosmos_framework/configs/simulation/config.py | 34 +- .../action_caption_attribute_adapter.py | 17 + .../generator/action/utils/camera_views.py | 112 +++ .../data/generator/action/utils/multiview.py | 136 +++ .../data/generator/action/utils/transforms.py | 111 ++- .../generator/action/utils/transforms_test.py | 249 ++++++ .../generator/action/utils/viewpoint_utils.py | 14 +- .../augmentors/interleaved_image_transform.py | 21 +- .../multi_reference_bbox_condition.py | 4 +- .../multi_reference_subject_layout.py | 490 ++++++++++ .../multi_reference_subject_layout_test.py | 249 ++++++ .../augmentors/multi_reference_transform.py | 22 +- .../augmentors/reasoner/bytes_to_media.py | 30 +- .../augmentors/reasoner/filter_seq_length.py | 28 +- .../reasoner/format_describe_anything.py | 8 +- .../augmentors/reasoner/prompt_format.py | 28 +- .../augmentors/reasoner/prompt_format_test.py | 35 + .../augmentors/reasoner/sample_data_stats.py | 3 +- .../augmentors/reasoner/timestamp.py | 5 +- .../timestamp_with_subject_tracking.py | 8 +- .../augmentors/reasoner/tokenize_data.py | 434 ++++++++- .../augmentors/reasoner/tokenize_data_test.py | 290 +++++- .../augmentors/target_first_frame_test.py | 48 + .../augmentors/text_transforms_for_video.py | 2 +- .../augmentors/torchcodec_callers_test.py | 2 + .../augmentors/video_editing_transform.py | 65 +- .../video_editing_transform_test.py | 79 ++ .../indexed_transfer_length_loader.py | 2 +- .../data/generator/joint_dataloader.py | 838 +++++++++--------- .../data/generator/lance/augment_timing.py | 25 +- .../data/generator/local_simulation/config.py | 2 +- .../generator/local_simulation/dataloader.py | 2 +- .../generator/multiview/camera_attributes.py | 1 + .../multiview/camera_attributes_test.py | 1 + .../generator/multiview/camera_geometry.py | 23 +- .../multiview/camera_pose_actions.py | 678 ++++++++++++++ .../generator/multiview/caption_format.py | 19 +- .../generator/processors/qwen3vl_processor.py | 126 ++- .../processors/qwen3vl_processor_test.py | 102 +++ .../data/generator/reasoner/collate_fn.py | 9 +- .../generator/reasoner/video_decoder_qwen.py | 153 +++- .../data/generator/sample_token_counter.py | 346 ++++++++ .../sequence_packing/autoregressive.py | 4 + .../generator/sequence_packing/packers.py | 4 +- .../sequence_packing/packers_test.py | 26 + .../generator/sequence_packing/sequence.py | 6 + .../data/generator/token_mix_control.py | 131 ++- .../data/generator/transfer_length_index.py | 4 +- .../common/distillation_export_test.py | 48 + cosmos_framework/model/_base.py | 7 + .../generator/algorithm/loss/cross_entropy.py | 175 +++- .../generator/algorithm/loss/flow_matching.py | 6 +- .../algorithm/loss/fp32_head_gradient_test.py | 139 +++ .../loss/fused_packed_cross_entropy_test.py | 60 ++ .../algorithm/loss/load_balancing_test.py | 187 +++- .../model/generator/context_parallel.py | 228 +++++ cosmos_framework/model/generator/hf_model.py | 104 ++- .../generator/mot/activation_marks_test.py | 4 +- .../model/generator/mot/causal_attention.py | 89 +- .../mot/causal_context_parallel_test.py | 16 +- .../mot/causal_cosmos3_vfm_network.py | 2 + .../generator/mot/context_parallel_test.py | 77 +- .../generator/mot/context_parallel_utils.py | 64 +- .../generator/mot/cosmos3_vfm_network.py | 522 ++++++----- .../model/generator/mot/flex_attention.py | 8 +- .../generator/mot/flex_attention_test.py | 311 ++++++- ...view_action_conditioning_attention_test.py | 244 +++++ .../generator/mot/multiview_attention.md | 21 + .../generator/mot/multiview_attention.py | 3 + .../mot/multiview_maskless_attention.py | 207 ++++- .../mot/multiview_maskless_exact_test.py | 221 ++++- .../generator/mot/parallelize_unified_mot.py | 8 +- .../mot/parallelize_unified_mot_test.py | 40 + .../generator/mot/three_way_attention_test.py | 16 +- .../model/generator/mot/unified_mot.py | 136 ++- .../model/generator/mot/unified_mot_test.py | 113 ++- .../generator/mot/vae_load_balance_test.py | 3 + .../model/generator/omni_mot_causal_model.py | 62 +- .../generator/omni_mot_causal_model_test.py | 103 ++- .../model/generator/omni_mot_model.py | 245 ++++- .../model/generator/parallelize_vlm.py | 77 +- .../generator/reasoner/audio/registry.py | 28 +- .../model/generator/reasoner/audio/utils.py | 38 +- .../generator/reasoner/audio/utils_test.py | 54 ++ .../qwen35_caption_accumulation_test.py | 72 +- .../configuration_qwen3_omni_audio.py | 14 +- .../qwen3_omni_audio/qwen3_omni_audio.py | 58 +- .../generator/reasoner/qwen3_vl_fused_ce.py | 64 ++ .../reasoner/qwen3_vl_moe/fixed_sum.py | 190 ++++ .../reasoner/qwen3_vl_moe/moe_test.py | 87 ++ .../model/generator/reasoner/qwen3_vl_rope.py | 41 + .../tokenizers/lidar/lidar_tokenizer_v1.py | 3 +- .../model/generator/utils/component_unipc.py | 100 +++ .../generator/utils/component_unipc_test.py | 95 ++ .../model/generator/utils/kv_cache.py | 2 +- .../model/generator/utils/kv_cache_test.py | 16 +- .../model/generator/utils/memory.py | 30 +- .../generator/utils/model_construction.py | 29 + .../model/generator/utils/moe_utils.py | 35 + .../model/generator/utils/multiview_ar.py | 5 + .../utils/rolling_kv/rolling_replay.py | 2 + .../utils/rolling_kv/rolling_transfer_ar.py | 6 + cosmos_framework/model/generator/vlm_model.py | 264 ++++-- cosmos_framework/scripts/train_experiment.py | 10 + cosmos_framework/trainer/__init__.py | 130 ++- cosmos_framework/trainer/distillation.py | 4 +- cosmos_framework/trainer/distillation_test.py | 2 +- cosmos_framework/utils/callback.py | 16 +- cosmos_framework/utils/config.py | 4 +- cosmos_framework/utils/distributed.py | 2 +- .../utils/easy_io/backends/__init__.py | 12 +- .../utils/easy_io/backends/registry_utils.py | 5 +- cosmos_framework/utils/easy_io/file_client.py | 6 +- .../utils/generator/attention_chunks.py | 28 + .../utils/generator/clusterpulse_preflight.py | 67 ++ .../utils/generator/cost_model/benchmark.py | 506 +++++++---- .../utils/generator/deterministic_rng.py | 19 +- .../utils/generator/fused_adam.py | 6 + .../utils/generator/hf_attention_cosmos.py | 107 ++- .../utils/generator/monkey_patch.py | 125 ++- .../utils/generator/parallelism.py | 12 +- .../utils/generator/reasoner/constant.py | 4 + .../utils/generator/reasoner/model_path.py | 29 + .../utils/generator/reasoner/true_packing.py | 34 +- .../generator/structured_conversation.py | 117 +++ .../generator/torchcodec_dependencies.py | 6 +- .../utils/reasoner/dcp_checkpointer.py | 31 +- 147 files changed, 9665 insertions(+), 1700 deletions(-) create mode 100644 cosmos_framework/data/generator/action/utils/camera_views.py create mode 100644 cosmos_framework/data/generator/action/utils/multiview.py create mode 100644 cosmos_framework/data/generator/augmentors/multi_reference_subject_layout.py create mode 100644 cosmos_framework/data/generator/augmentors/multi_reference_subject_layout_test.py create mode 100644 cosmos_framework/data/generator/augmentors/target_first_frame_test.py create mode 100644 cosmos_framework/data/generator/multiview/camera_pose_actions.py create mode 100644 cosmos_framework/data/generator/processors/qwen3vl_processor_test.py create mode 100644 cosmos_framework/data/generator/sample_token_counter.py create mode 100644 cosmos_framework/model/generator/algorithm/loss/fp32_head_gradient_test.py create mode 100644 cosmos_framework/model/generator/algorithm/loss/fused_packed_cross_entropy_test.py create mode 100644 cosmos_framework/model/generator/context_parallel.py create mode 100644 cosmos_framework/model/generator/mot/multiview_action_conditioning_attention_test.py create mode 100644 cosmos_framework/model/generator/reasoner/qwen3_vl_fused_ce.py create mode 100644 cosmos_framework/model/generator/reasoner/qwen3_vl_moe/fixed_sum.py create mode 100644 cosmos_framework/model/generator/reasoner/qwen3_vl_rope.py create mode 100644 cosmos_framework/model/generator/utils/component_unipc.py create mode 100644 cosmos_framework/model/generator/utils/component_unipc_test.py create mode 100644 cosmos_framework/model/generator/utils/model_construction.py create mode 100644 cosmos_framework/utils/generator/attention_chunks.py create mode 100644 cosmos_framework/utils/generator/clusterpulse_preflight.py create mode 100644 cosmos_framework/utils/generator/reasoner/model_path.py create mode 100644 cosmos_framework/utils/generator/structured_conversation.py diff --git a/.file_mapping.json b/.file_mapping.json index d740353be..588ab449f 100644 --- a/.file_mapping.json +++ b/.file_mapping.json @@ -1,7 +1,7 @@ { - "_source_commit": "a453217f601aae1150a016131f81ccfd89f5e3dc", - "_dest_commit": "60ad5b8d221de927d5bec80938d6d456dc968f13", - "_generated_at": "2026-10-14T15:32:11Z", + "_source_commit": "b3866eb30721bf68dce73380563235d39fa5a622-dirty", + "_dest_commit": "bf492fe2fda2108aae9b9561dc9fb9eec3fe2027", + "_generated_at": "2026-10-24T15:53:46Z", "files": { "imaginaire/__init__.py": "cosmos_framework/__init__.py", "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py", @@ -200,6 +200,8 @@ "projects/cosmos3/cosmos3/algorithm/loss/distillation_test.py": "cosmos_framework/model/generator/distillation/common_loss_test.py", "projects/cosmos3/cosmos3/algorithm/loss/flow_matching.py": "cosmos_framework/model/generator/algorithm/loss/flow_matching.py", "projects/cosmos3/cosmos3/algorithm/loss/flow_matching_test.py": "cosmos_framework/model/generator/algorithm/loss/flow_matching_test.py", + "projects/cosmos3/cosmos3/algorithm/loss/fp32_head_gradient_test.py": "cosmos_framework/model/generator/algorithm/loss/fp32_head_gradient_test.py", + "projects/cosmos3/cosmos3/algorithm/loss/fused_packed_cross_entropy_test.py": "cosmos_framework/model/generator/algorithm/loss/fused_packed_cross_entropy_test.py", "projects/cosmos3/cosmos3/algorithm/loss/load_balancing.py": "cosmos_framework/model/generator/algorithm/loss/load_balancing.py", "projects/cosmos3/cosmos3/algorithm/loss/load_balancing_test.py": "cosmos_framework/model/generator/algorithm/loss/load_balancing_test.py", "projects/cosmos3/cosmos3/algorithm/loss/modality_reduction.py": "cosmos_framework/model/generator/algorithm/loss/modality_reduction.py", @@ -298,7 +300,6 @@ "projects/cosmos3/cosmos3/configs/simulation/defaults/flex_attention.py": "cosmos_framework/configs/base/defaults/causal_flex_attention.py", "projects/cosmos3/cosmos3/configs/simulation/defaults/optimizer.py": "cosmos_framework/configs/base/defaults/distill_optimizer.py", "projects/cosmos3/cosmos3/configs/simulation/defaults/replay_attention.py": "cosmos_framework/configs/base/defaults/replay_attention.py", - "projects/cosmos3/cosmos3/configs/simulation/local.py": "cosmos_framework/configs/simulation/local.py", "projects/cosmos3/cosmos3/configs/simulation/methods/__init__.py": "cosmos_framework/configs/base/experiment/distillation/__init__.py", "projects/cosmos3/cosmos3/configs/simulation/methods/dmd2.py": "cosmos_framework/configs/base/experiment/distillation/dmd2_config.py", "projects/cosmos3/cosmos3/configs/simulation/methods/dmd2_test.py": "cosmos_framework/configs/base/experiment/distillation/dmd2_config_test.py", @@ -319,8 +320,10 @@ "projects/cosmos3/cosmos3/datasets/action/utils/action_spec.py": "cosmos_framework/data/generator/action/utils/action_spec.py", "projects/cosmos3/cosmos3/datasets/action/utils/camera_trajectory.py": "cosmos_framework/data/generator/action/utils/camera_trajectory.py", "projects/cosmos3/cosmos3/datasets/action/utils/camera_trajectory_test.py": "cosmos_framework/data/generator/action/utils/camera_trajectory_test.py", + "projects/cosmos3/cosmos3/datasets/action/utils/camera_views.py": "cosmos_framework/data/generator/action/utils/camera_views.py", "projects/cosmos3/cosmos3/datasets/action/utils/domain_utils.py": "cosmos_framework/data/generator/action/utils/domain_utils.py", "projects/cosmos3/cosmos3/datasets/action/utils/json_formatter.py": "cosmos_framework/data/generator/action/utils/json_formatter.py", + "projects/cosmos3/cosmos3/datasets/action/utils/multiview.py": "cosmos_framework/data/generator/action/utils/multiview.py", "projects/cosmos3/cosmos3/datasets/action/utils/pose_utils.py": "cosmos_framework/data/generator/action/utils/pose_utils.py", "projects/cosmos3/cosmos3/datasets/action/utils/pose_utils_test.py": "cosmos_framework/data/generator/action/utils/pose_utils_test.py", "projects/cosmos3/cosmos3/datasets/action/utils/transforms.py": "cosmos_framework/data/generator/action/utils/transforms.py", @@ -352,6 +355,8 @@ "projects/cosmos3/cosmos3/datasets/augmentors/merge_datadict.py": "cosmos_framework/data/generator/augmentors/merge_datadict.py", "projects/cosmos3/cosmos3/datasets/augmentors/multi_reference_bbox_condition.py": "cosmos_framework/data/generator/augmentors/multi_reference_bbox_condition.py", "projects/cosmos3/cosmos3/datasets/augmentors/multi_reference_bbox_condition_test.py": "cosmos_framework/data/generator/augmentors/multi_reference_bbox_condition_test.py", + "projects/cosmos3/cosmos3/datasets/augmentors/multi_reference_subject_layout.py": "cosmos_framework/data/generator/augmentors/multi_reference_subject_layout.py", + "projects/cosmos3/cosmos3/datasets/augmentors/multi_reference_subject_layout_test.py": "cosmos_framework/data/generator/augmentors/multi_reference_subject_layout_test.py", "projects/cosmos3/cosmos3/datasets/augmentors/multi_reference_transform.py": "cosmos_framework/data/generator/augmentors/multi_reference_transform.py", "projects/cosmos3/cosmos3/datasets/augmentors/pkl_to_media.py": "cosmos_framework/data/generator/augmentors/pkl_to_media.py", "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/__init__.py": "cosmos_framework/data/generator/augmentors/reasoner/__init__.py", @@ -381,6 +386,7 @@ "projects/cosmos3/cosmos3/datasets/augmentors/sequence_plan.py": "cosmos_framework/data/generator/augmentors/sequence_plan.py", "projects/cosmos3/cosmos3/datasets/augmentors/sound_sequence_plan.py": "cosmos_framework/data/generator/augmentors/sound_sequence_plan.py", "projects/cosmos3/cosmos3/datasets/augmentors/sound_sequence_plan_test.py": "cosmos_framework/data/generator/augmentors/sound_sequence_plan_test.py", + "projects/cosmos3/cosmos3/datasets/augmentors/target_first_frame_test.py": "cosmos_framework/data/generator/augmentors/target_first_frame_test.py", "projects/cosmos3/cosmos3/datasets/augmentors/text_tokenizer.py": "cosmos_framework/data/generator/augmentors/text_tokenizer.py", "projects/cosmos3/cosmos3/datasets/augmentors/text_tokenizer_test.py": "cosmos_framework/data/generator/augmentors/text_tokenizer_test.py", "projects/cosmos3/cosmos3/datasets/augmentors/text_transforms_for_image.py": "cosmos_framework/data/generator/augmentors/text_transforms_for_image.py", @@ -412,7 +418,6 @@ "projects/cosmos3/cosmos3/datasets/local_simulation/config.py": "cosmos_framework/data/generator/local_simulation/config.py", "projects/cosmos3/cosmos3/datasets/local_simulation/core.py": "cosmos_framework/data/generator/local_simulation/core.py", "projects/cosmos3/cosmos3/datasets/local_simulation/dataloader.py": "cosmos_framework/data/generator/local_simulation/dataloader.py", - "projects/cosmos3/cosmos3/datasets/local_simulation/dataset.py": "cosmos_framework/data/generator/local_simulation/dataset.py", "projects/cosmos3/cosmos3/datasets/local_simulation/resume.py": "cosmos_framework/data/generator/local_simulation/resume.py", "projects/cosmos3/cosmos3/datasets/local_simulation/transfer_bundle.py": "cosmos_framework/data/generator/local_simulation/transfer_bundle.py", "projects/cosmos3/cosmos3/datasets/multiview/camera_attributes.py": "cosmos_framework/data/generator/multiview/camera_attributes.py", @@ -420,6 +425,7 @@ "projects/cosmos3/cosmos3/datasets/multiview/camera_geometry.py": "cosmos_framework/data/generator/multiview/camera_geometry.py", "projects/cosmos3/cosmos3/datasets/multiview/camera_geometry_contract.py": "cosmos_framework/data/generator/multiview/camera_geometry_contract.py", "projects/cosmos3/cosmos3/datasets/multiview/camera_geometry_native.py": "cosmos_framework/data/generator/multiview/camera_geometry_native.py", + "projects/cosmos3/cosmos3/datasets/multiview/camera_pose_actions.py": "cosmos_framework/data/generator/multiview/camera_pose_actions.py", "projects/cosmos3/cosmos3/datasets/multiview/caption_chunks.py": "cosmos_framework/data/generator/multiview/caption_chunks.py", "projects/cosmos3/cosmos3/datasets/multiview/caption_format.py": "cosmos_framework/data/generator/multiview/caption_format.py", "projects/cosmos3/cosmos3/datasets/packing_iterable_dataset.py": "cosmos_framework/data/generator/packing_iterable_dataset.py", @@ -427,6 +433,7 @@ "projects/cosmos3/cosmos3/datasets/reasoner/collate_fn.py": "cosmos_framework/data/generator/reasoner/collate_fn.py", "projects/cosmos3/cosmos3/datasets/reasoner/video_decoder_qwen.py": "cosmos_framework/data/generator/reasoner/video_decoder_qwen.py", "projects/cosmos3/cosmos3/datasets/recipe_database_api.py": "cosmos_framework/data/generator/recipe_database_api.py", + "projects/cosmos3/cosmos3/datasets/sample_token_counter.py": "cosmos_framework/data/generator/sample_token_counter.py", "projects/cosmos3/cosmos3/datasets/sound_data_utils.py": "cosmos_framework/data/generator/sound_data_utils.py", "projects/cosmos3/cosmos3/datasets/token_mix_control.py": "cosmos_framework/data/generator/token_mix_control.py", "projects/cosmos3/cosmos3/datasets/transfer_augmentors.py": "cosmos_framework/data/generator/transfer_augmentors.py", @@ -448,6 +455,7 @@ "projects/cosmos3/cosmos3/diffusion/samplers/utils_test.py": "cosmos_framework/model/generator/diffusion/samplers/utils_test.py", "projects/cosmos3/cosmos3/models/attention_io_layout.py": "cosmos_framework/model/generator/attention_io_layout.py", "projects/cosmos3/cosmos3/models/consistency_model/tf_dcm.py": "cosmos_framework/model/generator/consistency_model/tf_dcm.py", + "projects/cosmos3/cosmos3/models/context_parallel.py": "cosmos_framework/model/generator/context_parallel.py", "projects/cosmos3/cosmos3/models/distribution_matching/__init__.py": "cosmos_framework/model/generator/distillation/__init__.py", "projects/cosmos3/cosmos3/models/distribution_matching/dmd2_rf.py": "cosmos_framework/model/generator/distillation/dmd2_rf.py", "projects/cosmos3/cosmos3/models/distribution_matching/dmd2_rf_test.py": "cosmos_framework/model/generator/distillation/dmd2_rf_test.py", @@ -457,11 +465,6 @@ "projects/cosmos3/cosmos3/models/distribution_matching/trainer.py": "cosmos_framework/trainer/distillation.py", "projects/cosmos3/cosmos3/models/distribution_matching/trainer_test.py": "cosmos_framework/trainer/distillation_test.py", "projects/cosmos3/cosmos3/models/hf_model.py": "cosmos_framework/model/generator/hf_model.py", - "projects/cosmos3/cosmos3/models/moe/deep_gemm.py": "cosmos_framework/model/generator/moe/deep_gemm.py", - "projects/cosmos3/cosmos3/models/moe/flash_infer.py": "cosmos_framework/model/generator/moe/flash_infer.py", - "projects/cosmos3/cosmos3/models/moe/interface.py": "cosmos_framework/model/generator/moe/interface.py", - "projects/cosmos3/cosmos3/models/moe/native.py": "cosmos_framework/model/generator/moe/native.py", - "projects/cosmos3/cosmos3/models/moe/native_kernels.py": "cosmos_framework/model/generator/moe/native_kernels.py", "projects/cosmos3/cosmos3/models/mot/__init__.py": "cosmos_framework/model/generator/mot/__init__.py", "projects/cosmos3/cosmos3/models/mot/action_io_projector.py": "cosmos_framework/model/generator/mot/action_io_projector.py", "projects/cosmos3/cosmos3/models/mot/action_io_projector_test.py": "cosmos_framework/model/generator/mot/action_io_projector_test.py", @@ -501,6 +504,7 @@ "projects/cosmos3/cosmos3/models/mot/modeling_utils.py": "cosmos_framework/model/generator/mot/modeling_utils.py", "projects/cosmos3/cosmos3/models/mot/modeling_utils_test.py": "cosmos_framework/model/generator/mot/modeling_utils_test.py", "projects/cosmos3/cosmos3/models/mot/multiview_action_attention.py": "cosmos_framework/model/generator/mot/multiview_action_attention.py", + "projects/cosmos3/cosmos3/models/mot/multiview_action_conditioning_attention_test.py": "cosmos_framework/model/generator/mot/multiview_action_conditioning_attention_test.py", "projects/cosmos3/cosmos3/models/mot/multiview_attention.md": "cosmos_framework/model/generator/mot/multiview_attention.md", "projects/cosmos3/cosmos3/models/mot/multiview_attention.py": "cosmos_framework/model/generator/mot/multiview_attention.py", "projects/cosmos3/cosmos3/models/mot/multiview_attention_test.py": "cosmos_framework/model/generator/mot/multiview_attention_test.py", @@ -578,25 +582,31 @@ "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl/qwen3_vl.py": "cosmos_framework/model/generator/reasoner/qwen3_vl/qwen3_vl.py", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl/utils.py": "cosmos_framework/model/generator/reasoner/qwen3_vl/utils.py", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl/video_processing_qwen3_vl.py": "cosmos_framework/model/generator/reasoner/qwen3_vl/video_processing_qwen3_vl.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_fused_ce.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_fused_ce.py", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/__init__.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/__init__.py", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/configs/__init__.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/__init__.py", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/configuration_qwen3_vl_moe.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configuration_qwen3_vl_moe.py", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/cosine_router_test.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/cosine_router_test.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/fixed_sum.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/fixed_sum.py", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/moe.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe.py", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/moe_bench.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_bench.py", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/moe_kernels.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_kernels.py", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/moe_test.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_test.py", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/qwen3_vl_moe.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/qwen3_vl_moe.py", "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/shared_expert_test.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/shared_expert_test.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_rope.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_rope.py", "projects/cosmos3/cosmos3/models/sensor_encoder.py": "cosmos_framework/model/generator/sensor_encoder.py", "projects/cosmos3/cosmos3/models/utils/__init__.py": "cosmos_framework/model/generator/utils/__init__.py", + "projects/cosmos3/cosmos3/models/utils/action_stream.py": "cosmos_framework/model/generator/utils/action_stream.py", "projects/cosmos3/cosmos3/models/utils/batch_normalization.py": "cosmos_framework/model/generator/utils/batch_normalization.py", "projects/cosmos3/cosmos3/models/utils/camera_rays.py": "cosmos_framework/model/generator/utils/camera_rays.py", "projects/cosmos3/cosmos3/models/utils/camera_rays_test.py": "cosmos_framework/model/generator/utils/camera_rays_test.py", "projects/cosmos3/cosmos3/models/utils/camera_relative_pose.py": "cosmos_framework/model/generator/utils/camera_relative_pose.py", "projects/cosmos3/cosmos3/models/utils/camera_relative_pose_test.py": "cosmos_framework/model/generator/utils/camera_relative_pose_test.py", + "projects/cosmos3/cosmos3/models/utils/component_unipc.py": "cosmos_framework/model/generator/utils/component_unipc.py", + "projects/cosmos3/cosmos3/models/utils/component_unipc_test.py": "cosmos_framework/model/generator/utils/component_unipc_test.py", "projects/cosmos3/cosmos3/models/utils/data_and_condition.py": "cosmos_framework/model/generator/utils/data_and_condition.py", "projects/cosmos3/cosmos3/models/utils/data_and_condition_test.py": "cosmos_framework/model/generator/utils/data_and_condition_test.py", "projects/cosmos3/cosmos3/models/utils/frustum_retrieval.py": "cosmos_framework/model/generator/utils/frustum_retrieval.py", @@ -610,9 +620,11 @@ "projects/cosmos3/cosmos3/models/utils/kv_storage_backend_test.py": "cosmos_framework/model/generator/utils/kv_storage_backend_test.py", "projects/cosmos3/cosmos3/models/utils/load_balancing_stats.py": "cosmos_framework/model/generator/utils/load_balancing_stats.py", "projects/cosmos3/cosmos3/models/utils/memory.py": "cosmos_framework/model/generator/utils/memory.py", + "projects/cosmos3/cosmos3/models/utils/model_construction.py": "cosmos_framework/model/generator/utils/model_construction.py", "projects/cosmos3/cosmos3/models/utils/moe_utils.py": "cosmos_framework/model/generator/utils/moe_utils.py", "projects/cosmos3/cosmos3/models/utils/multiview_ar.py": "cosmos_framework/model/generator/utils/multiview_ar.py", - "projects/cosmos3/cosmos3/models/utils/multiview_test_utils.py": "cosmos_framework/model/generator/utils/multiview_test_utils.py", + "projects/cosmos3/cosmos3/models/utils/numerics.py": "cosmos_framework/model/generator/utils/numerics.py", + "projects/cosmos3/cosmos3/models/utils/numerics_test.py": "cosmos_framework/model/generator/utils/numerics_test.py", "projects/cosmos3/cosmos3/models/utils/nvfp4.py": "cosmos_framework/model/generator/utils/nvfp4.py", "projects/cosmos3/cosmos3/models/utils/nvfp4_test.py": "cosmos_framework/model/generator/utils/nvfp4_test.py", "projects/cosmos3/cosmos3/models/utils/prope_cross_view_poses_test.py": "cosmos_framework/model/generator/utils/prope_cross_view_poses_test.py", @@ -649,6 +661,7 @@ "projects/cosmos3/cosmos3/processors/qwen3vl_nemo_chat_processor.py": "cosmos_framework/data/generator/processors/qwen3vl_nemo_chat_processor.py", "projects/cosmos3/cosmos3/processors/qwen3vl_nemo_chat_processor_test.py": "cosmos_framework/data/generator/processors/qwen3vl_nemo_chat_processor_test.py", "projects/cosmos3/cosmos3/processors/qwen3vl_processor.py": "cosmos_framework/data/generator/processors/qwen3vl_processor.py", + "projects/cosmos3/cosmos3/processors/qwen3vl_processor_test.py": "cosmos_framework/data/generator/processors/qwen3vl_processor_test.py", "projects/cosmos3/cosmos3/processors/source_video_timing_test.py": "cosmos_framework/data/generator/processors/source_video_timing_test.py", "projects/cosmos3/cosmos3/scripts/multiview_auto/multiview_collage.py": "cosmos_framework/scripts/multiview_collage.py", "projects/cosmos3/cosmos3/scripts/multiview_auto/multiview_collage_test.py": "cosmos_framework/scripts/multiview_collage_test.py", @@ -697,7 +710,9 @@ "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/activation_offloading.py": "cosmos_framework/utils/generator/activation_offloading.py", + "projects/cosmos3/cosmos3/utils/attention_chunks.py": "cosmos_framework/utils/generator/attention_chunks.py", "projects/cosmos3/cosmos3/utils/aux_optimizer_utils.py": "cosmos_framework/utils/generator/aux_optimizer_utils.py", + "projects/cosmos3/cosmos3/utils/clusterpulse_preflight.py": "cosmos_framework/utils/generator/clusterpulse_preflight.py", "projects/cosmos3/cosmos3/utils/data_batch.py": "cosmos_framework/utils/generator/data_batch.py", "projects/cosmos3/cosmos3/utils/data_utils.py": "cosmos_framework/utils/generator/data_utils.py", "projects/cosmos3/cosmos3/utils/deterministic_rng.py": "cosmos_framework/utils/generator/deterministic_rng.py", @@ -727,12 +742,15 @@ "projects/cosmos3/cosmos3/utils/reasoner/__init__.py": "cosmos_framework/utils/generator/reasoner/__init__.py", "projects/cosmos3/cosmos3/utils/reasoner/constant.py": "cosmos_framework/utils/generator/reasoner/constant.py", "projects/cosmos3/cosmos3/utils/reasoner/flop_calculator.py": "cosmos_framework/utils/generator/reasoner/flop_calculator.py", + "projects/cosmos3/cosmos3/utils/reasoner/model_path.py": "cosmos_framework/utils/generator/reasoner/model_path.py", "projects/cosmos3/cosmos3/utils/reasoner/packing_mbs_advisor.py": "cosmos_framework/utils/generator/reasoner/packing_mbs_advisor.py", "projects/cosmos3/cosmos3/utils/reasoner/pretrained_models_downloader.py": "cosmos_framework/utils/generator/reasoner/pretrained_models_downloader.py", "projects/cosmos3/cosmos3/utils/reasoner/pretrained_models_downloader_test.py": "cosmos_framework/utils/generator/reasoner/pretrained_models_downloader_test.py", "projects/cosmos3/cosmos3/utils/reasoner/true_packing.py": "cosmos_framework/utils/generator/reasoner/true_packing.py", "projects/cosmos3/cosmos3/utils/source_video_timing.py": "cosmos_framework/utils/generator/source_video_timing.py", "projects/cosmos3/cosmos3/utils/spatial_patch.py": "cosmos_framework/utils/generator/spatial_patch.py", + "projects/cosmos3/cosmos3/utils/structured_conversation.py": "cosmos_framework/utils/generator/structured_conversation.py", + "projects/cosmos3/cosmos3/utils/torchcodec_dependencies.py": "cosmos_framework/utils/generator/torchcodec_dependencies.py", "projects/cosmos3/cosmos3/utils/video_frame_sampling.py": "cosmos_framework/utils/generator/video_frame_sampling.py", "projects/cosmos3/cosmos3/utils/video_preprocess.py": "cosmos_framework/utils/generator/video_preprocess.py", "projects/cosmos3/cosmos3/utils/video_source_metadata.py": "cosmos_framework/utils/generator/video_source_metadata.py", diff --git a/cosmos_framework/auxiliary/guardrail/common/presets.py b/cosmos_framework/auxiliary/guardrail/common/presets.py index 4dc00ed40..8c4e2aee6 100644 --- a/cosmos_framework/auxiliary/guardrail/common/presets.py +++ b/cosmos_framework/auxiliary/guardrail/common/presets.py @@ -9,12 +9,15 @@ from cosmos_framework.utils import log -def create_text_guardrail_runner(offload_model_to_cpu: bool = False) -> GuardrailRunner: +def create_text_guardrail_runner( + offload_model_to_cpu: bool = False, *, raise_on_error: bool = False +) -> GuardrailRunner: """Create the text guardrail runner.""" + strict_options: dict[str, bool] = {"raise_on_error": True} if raise_on_error else {} return GuardrailRunner( safety_models=[ Blocklist(), - Qwen3Guard(offload_model_to_cpu=offload_model_to_cpu), + Qwen3Guard(offload_model_to_cpu=offload_model_to_cpu, **strict_options), ] ) diff --git a/cosmos_framework/auxiliary/guardrail/qwen3guard/qwen3guard.py b/cosmos_framework/auxiliary/guardrail/qwen3guard/qwen3guard.py index b165db153..dd3bc6063 100644 --- a/cosmos_framework/auxiliary/guardrail/qwen3guard/qwen3guard.py +++ b/cosmos_framework/auxiliary/guardrail/qwen3guard/qwen3guard.py @@ -31,19 +31,23 @@ class Qwen3Guard(ContentSafetyGuardrail): model_id: str model: Any tokenizer: Any + raise_on_error: bool def __init__( self, offload_model_to_cpu: bool = True, + raise_on_error: bool = False, ) -> None: """Load Qwen3Guard for text filtering safety checks. Args: offload_model_to_cpu (bool, optional): Whether to offload the model to CPU. Defaults to True. + raise_on_error (bool): Propagate runtime errors for strict evaluation. Defaults to False. """ from transformers import AutoModelForCausalLM, AutoTokenizer self.offload_model = offload_model_to_cpu + self.raise_on_error = raise_on_error self.dtype = torch.bfloat16 self.model_id = "Qwen/Qwen3Guard-Gen-0.6B" @@ -95,6 +99,8 @@ def is_safe(self, prompt: str) -> tuple[bool, str]: try: return self.extract_label_and_categories(prompt) except Exception as e: + if self.raise_on_error: + raise log.error(f"Unexpected error occurred when running Qwen3Guard guardrail: {e}") return True, "Unexpected error occurred when running Qwen3Guard guardrail." diff --git a/cosmos_framework/callbacks/data_stats_writer.py b/cosmos_framework/callbacks/data_stats_writer.py index 7d815f608..f52fabde0 100644 --- a/cosmos_framework/callbacks/data_stats_writer.py +++ b/cosmos_framework/callbacks/data_stats_writer.py @@ -52,7 +52,7 @@ def arrow(self) -> pa.Field: _FIELDS = ( _Field("phase", pa.string(), 'Data split that produced the row: "train", "val", or "data_only".'), - _Field("step", pa.int64(), "Completed optimizer step associated with the sample."), + _Field("step", pa.int64(), "Completed training iteration or data-loading step associated with the sample."), _Field("rank", pa.int32(), "Distributed process rank that consumed the sample."), _Field("sample_rank", pa.int32(), "Rank recorded by the dataset distributor."), _Field("worker_id", pa.int32(), "Dataloader worker that selected the sample."), @@ -111,7 +111,7 @@ def arrow(self) -> pa.Field: _Field("context_length", pa.int32(), "Configured model context length."), _Field("dataset_weight", pa.float64(), "Effective normalized sampling share of the dataset."), _Field("dataset_num_urls", pa.int64(), "Dataset URL count after recipe expansion."), - _Field("total_training_steps", pa.int64(), "Configured maximum optimizer iterations."), + _Field("total_training_steps", pa.int64(), "Configured maximum training or data-loading iterations."), ) DATA_STATS_SCHEMA = pa.schema( @@ -290,6 +290,9 @@ class DataStatsWriterCallback(Callback): training continues. The next save, the end of training, and resume each wait for that merge and stop the job if it cannot succeed, so rank files never accumulate unnoticed. Rows recorded after the final save, which can only come from the last validation, are not saved. + + With ``trainer.skip_model_update``, train rows retain sample metadata and trainable-token counts; + CE and objective fields are null because no model forward pass ran. """ def __init__( @@ -349,6 +352,8 @@ def _record( return if not self._ready: raise RuntimeError("DataStatsWriterCallback was not initialized") + if int(data_batch.get("context_parallel_size", 1)) > 1: + raise NotImplementedError("Disable per-sample data-stats writing for CP; shard aggregation is required") metadata = data_batch.get("data_stats") if not isinstance(metadata, list) or not all(isinstance(sample_metadata, dict) for sample_metadata in metadata): raise TypeError( @@ -357,6 +362,8 @@ def _record( ) if data_batch.get("true_packing") or "cu_seq_lens_q" in data_batch: raise NotImplementedError("Per-sample data-stats logging currently requires padded batches") + if phase == "train" and getattr(self.config.trainer, "skip_model_update", False): + output_batch = None missing = [key for key in _LOSS_KEYS if output_batch is not None and key not in output_batch] if missing: raise KeyError(f"Model output is missing data-stats loss fields: {missing}") @@ -374,7 +381,7 @@ def _record( batch_values.pop("dataset_num_urls") _check_sample(keys[index], sample_metadata, batch_values) if output_batch is None: - labels = data_batch["labels"] + labels = data_batch["labels"] # [B,T] ignore_indices = _values(data_batch["ignore_index"], size) valid_counts = [int((labels[index, 1:] != ignore_indices[index]).sum()) for index in range(size)] token_sums = numerators = weights = [None] * size diff --git a/cosmos_framework/callbacks/dataloader_state.py b/cosmos_framework/callbacks/dataloader_state.py index af1977ba2..5d18b751b 100644 --- a/cosmos_framework/callbacks/dataloader_state.py +++ b/cosmos_framework/callbacks/dataloader_state.py @@ -55,7 +55,7 @@ def __init__( self.lance_state: dict[int, dict[str, Any]] = {} self.val_lance_state: dict[int, dict[str, Any]] = {} self.val_lance_batches: int = 0 - self._pending_lance_state: dict[str, Any] | None = None + self._pending_lance_state: dict[str, Any] | list[dict[str, Any]] | None = None self._pending_val_lance_state: dict[str, Any] | None = None self.verbose = True @@ -75,7 +75,12 @@ def on_training_step_batch_start( return lance_state = data_batch.pop(LANCE_VLM_RESUME_STATE_KEY, None) if lance_state is not None and not isinstance(lance_state, dict): - raise TypeError(f"{LANCE_VLM_RESUME_STATE_KEY} must be a dict, got {type(lance_state).__name__}.") + if ( + not isinstance(lance_state, list) + or not lance_state + or not all(isinstance(x, dict) for x in lance_state) + ): + raise TypeError(f"{LANCE_VLM_RESUME_STATE_KEY} must be a dict or a non-empty list of cursor dicts.") self._pending_lance_state = lance_state def _update_lance_state(self, update: dict[str, Any], split: str = "train") -> None: @@ -168,7 +173,9 @@ def on_training_step_batch_end( ) -> None: if self.distributor_type == "no_replace": if self._pending_lance_state is not None: - self._update_lance_state(self._pending_lance_state) + updates = self._pending_lance_state + for update in updates if isinstance(updates, list) else [updates]: + self._update_lance_state(update) self._pending_lance_state = None else: self._update_state_from_batch(self.state, data_batch) diff --git a/cosmos_framework/callbacks/every_n_draw_sample.py b/cosmos_framework/callbacks/every_n_draw_sample.py index 314613ea2..6cf1f288c 100644 --- a/cosmos_framework/callbacks/every_n_draw_sample.py +++ b/cosmos_framework/callbacks/every_n_draw_sample.py @@ -3,8 +3,8 @@ import math import os -from collections.abc import Callable -from contextlib import nullcontext +from collections.abc import Callable, Iterator +from contextlib import contextmanager, nullcontext from dataclasses import dataclass from functools import partial from typing import Any, Literal, NamedTuple @@ -30,6 +30,7 @@ TEXT_SYSTEM_PROMPT_KEY, ) from cosmos_framework.model.generator.mot.context_parallel_utils import broadcast_context_parallel_object +from cosmos_framework.model.generator.utils.moe_utils import preserve_moe_learning_accumulators from cosmos_framework.utils.generator.data_utils import slice_data_batch from cosmos_framework.utils.generator.multiview import ( decode_multiview_latent_per_view, @@ -50,6 +51,17 @@ class WandbAnimation(NamedTuple): WandbMedia = str | WandbAnimation | list[wandb.Object3D] | dict[str, "WandbMedia"] +@contextmanager +def _sampling_observation_scope(model: Any) -> Iterator[None]: + """Keep diagnostic sampling from contributing to the next router update.""" + net = getattr(model, "net", None) + if not isinstance(net, torch.nn.Module): + yield + return + with preserve_moe_learning_accumulators(net): + yield + + def resize_image(image: torch.Tensor, size: int = 1024) -> torch.Tensor: """ Resize the image to the given size. This is done so that wandb can display the image correctly. @@ -947,7 +959,9 @@ def every_n_impl( ) log.debug("entering, every_n_impl", rank0_only=False) - with context(): + # Keep this scope outside ``ema_scope`` so its ``finally`` restores learning + # accumulators after EMA weights have been swapped back, including on exceptions. + with _sampling_observation_scope(model), context(): if self.do_x0_prediction: log.debug("entering, x0_pred", rank0_only=False) x0_img_fp, mse_loss, sigmas = self.x0_pred( diff --git a/cosmos_framework/callbacks/sampled_media_recorder.py b/cosmos_framework/callbacks/sampled_media_recorder.py index ab82ad0e5..4a1b7a856 100644 --- a/cosmos_framework/callbacks/sampled_media_recorder.py +++ b/cosmos_framework/callbacks/sampled_media_recorder.py @@ -181,7 +181,7 @@ def _extract_records(self, data_batch: dict[str, Any], iteration: int, rank: int count = len(sample_ids) media_urls = _as_list(data_batch.get("__url__"), count) dataset_names = _as_list(data_batch.get("dataset_name"), count) - source_streams = _as_list(data_batch.get("transfer_source_stream"), count) + source_streams = _as_list(data_batch.get("task_family"), count) source_names = _as_list(data_batch.get("source_dataset_name"), count) captions = _as_optional_string_list(data_batch.get("ai_caption"), count) if self.record_caption else [] source_ids = _as_list(data_batch.get("source_id"), count) diff --git a/cosmos_framework/callbacks/sampled_media_recorder_test.py b/cosmos_framework/callbacks/sampled_media_recorder_test.py index b13d03f4b..3ff6343f2 100644 --- a/cosmos_framework/callbacks/sampled_media_recorder_test.py +++ b/cosmos_framework/callbacks/sampled_media_recorder_test.py @@ -59,12 +59,12 @@ def test_extract_records_from_consumed_image_batch(monkeypatch: pytest.MonkeyPat assert all(record["rank"] == 3 for record in records) -def test_extract_records_preserves_mixed_transfer_source_stream() -> None: +def test_extract_records_preserves_token_mixed_source_stream() -> None: callback = SampledMediaRecorder(enabled=True, output_uri="/tmp/samples.lance") batch = { "__key__": ["sample-a"], "dataset_name": ["transfer_rgb"], - "transfer_source_stream": ["transfer_depth_single_view_720"], + "task_family": ["transfer_depth_single_view_720"], } records = callback._extract_records(batch, iteration=7, rank=3) diff --git a/cosmos_framework/callbacks/tokens_per_sec.py b/cosmos_framework/callbacks/tokens_per_sec.py index 6cf522405..e3b1e6589 100644 --- a/cosmos_framework/callbacks/tokens_per_sec.py +++ b/cosmos_framework/callbacks/tokens_per_sec.py @@ -114,6 +114,7 @@ def __init__( self._fwd_bwd_flop_coeff: float | None = fwd_bwd_flop_coeff self._flop_coeff_resolved: float | None = None # cached auto-resolved coeff self._num_params: int | None = None # lazily filled from model.parameters() (global/DTensor) + self._context_parallel_size: int = 1 self._hit_counter: int = 0 self._tokens_in_window: int = 0 self._useful_tokens_in_window: int = 0 @@ -179,6 +180,7 @@ def on_training_step_batch_end( if self._hit_counter < self.hit_thres: return + self._context_parallel_size = int(data_batch.get("context_parallel_size", 1)) ids = data_batch.get(self.length_key) if ids is not None: # numel() is the physical token count. shape[0] counts rows, while a true-packed @@ -338,7 +340,8 @@ def _synchronize_window(self, elapsed: float) -> tuple[float, int, float, float, if dist.is_available() and dist.is_initialized() and dist.get_backend() == "nccl" else torch.device("cpu") ) - local_useful_tps = self._useful_tokens_in_window / elapsed if elapsed > 0 else 0.0 + cp_size = self._context_parallel_size + local_useful_tps = self._useful_tokens_in_window / cp_size / elapsed if elapsed > 0 else 0.0 local_predicted_microbatch_ms = ( self._predicted_runtime_ms_in_window / self._predicted_steps_in_window if self._predicted_steps_in_window @@ -371,6 +374,10 @@ def _synchronize_window(self, elapsed: float) -> tuple[float, int, float, float, dtype=torch.float64, device=reduce_device, ) + # The callback sees complete inputs replicated across CP ranks. Deduplicate + # token/sample/work totals, retaining physical-rank timing and step counts. + # Divide in FP64 before the global sum so small supervised counts stay exact. + sums[[0, 1, 2, 6, 7, 8, 14, 15]] /= cp_size peak_allocated_gb = torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0 peak_reserved_gb = torch.cuda.max_memory_reserved() / 1e9 if torch.cuda.is_available() else 0.0 maxima = torch.tensor( # [N_MAX] diff --git a/cosmos_framework/callbacks/wall_clock_checkpoint.py b/cosmos_framework/callbacks/wall_clock_checkpoint.py index 6b9b3c290..57efdb20d 100644 --- a/cosmos_framework/callbacks/wall_clock_checkpoint.py +++ b/cosmos_framework/callbacks/wall_clock_checkpoint.py @@ -263,7 +263,7 @@ def __init__( # hook has since described different optimizer/scheduler components. self._pending_dcp_components: dict[int, tuple[str, ...]] = {} - # Captured from on_before_optimizer_step so we can call checkpointer.save(). + # Captured before optimizer steps, or at train start for dataloading-only runs, to call checkpointer.save(). self._optimizer: torch.optim.Optimizer | None = None self._scheduler: torch.optim.lr_scheduler.LRScheduler | None = None self._grad_scaler: torch.amp.GradScaler | None = None @@ -309,6 +309,10 @@ def on_train_start(self, model: ImaginaireModel, iteration: int = 0) -> None: if not self.enabled: log.info(f"[WallClockCheckpoint] Disabled; set {INTERVAL_MINUTES_ENV_VAR} or interval_minutes to enable.") return + if getattr(self.config.trainer, "skip_model_update", False): + self._optimizer = self.trainer.optimizer + self._scheduler = self.trainer.scheduler + self._grad_scaler = self.trainer.grad_scaler log.info( f"[WallClockCheckpoint] Enabled: saving at least every {self._interval_seconds / 60:.1f} minutes, " f"retaining the last {self._keep_last} wall-clock checkpoints." diff --git a/cosmos_framework/checkpoint/base.py b/cosmos_framework/checkpoint/base.py index 1d607c8df..148ed3dc1 100644 --- a/cosmos_framework/checkpoint/base.py +++ b/cosmos_framework/checkpoint/base.py @@ -8,11 +8,11 @@ import torch +from cosmos_framework.utils.config import CheckpointConfig, JobConfig, ObjectStoreConfig +from cosmos_framework.utils.flags import INTERNAL from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.utils import callback -from cosmos_framework.utils.config import CheckpointConfig, JobConfig, ObjectStoreConfig from cosmos_framework.utils.easy_io import easy_io -from cosmos_framework.utils.flags import INTERNAL @dataclass(frozen=True, slots=True) diff --git a/cosmos_framework/checkpoint/dcp.py b/cosmos_framework/checkpoint/dcp.py index 7f2b5a861..5abaaef88 100644 --- a/cosmos_framework/checkpoint/dcp.py +++ b/cosmos_framework/checkpoint/dcp.py @@ -80,9 +80,9 @@ ) from cosmos_framework.checkpoint.base import AbstractCheckpointer, CheckpointLoadSource 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 diff --git a/cosmos_framework/checkpoint/dcp_distill.py b/cosmos_framework/checkpoint/dcp_distill.py index 39f9d8e5a..c8835fc51 100644 --- a/cosmos_framework/checkpoint/dcp_distill.py +++ b/cosmos_framework/checkpoint/dcp_distill.py @@ -27,6 +27,9 @@ from torch.distributed.checkpoint.stateful import Stateful from torch.nn.modules.module import _IncompatibleKeys +from cosmos_framework.model._base import ImaginaireModel +from cosmos_framework.utils import log, misc +from cosmos_framework.utils.easy_io import easy_io from cosmos_framework.checkpoint.dcp import ( AsyncMode, CustomLoadPlanner, @@ -36,10 +39,7 @@ DistributedCheckpointer as _DistributedCheckpointer, ) from cosmos_framework.checkpoint.dcp import ModelWrapper as VFMModelWrapper -from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.model.generator.distillation.optimizer import OptimizerContainerLike, is_optimizer_container -from cosmos_framework.utils import log, misc -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 __all__: tuple[str, ...] = ( diff --git a/cosmos_framework/configs/base/defaults/model_config.py b/cosmos_framework/configs/base/defaults/model_config.py index 11b6f5f4a..a9fb57308 100644 --- a/cosmos_framework/configs/base/defaults/model_config.py +++ b/cosmos_framework/configs/base/defaults/model_config.py @@ -340,6 +340,8 @@ class OmniMoTModelConfig: # None preserves Auto's separate action controls and control/target item roles. # Opt-in batches use separate ragged RGB and full frame-rate action blocks; # single-view batches in the same run keep their original dense packing. + # This config controls packing/view codes; attention uses the resulting item + # layout and the shared multiview_attention settings above. multiview_action_conditioning: MultiviewActionConditioningConfig | None = None # Per-layer NATTEN parameters diff --git a/cosmos_framework/configs/base/defaults/multiview_attention.py b/cosmos_framework/configs/base/defaults/multiview_attention.py index 3d6b99ef4..be85793d7 100644 --- a/cosmos_framework/configs/base/defaults/multiview_attention.py +++ b/cosmos_framework/configs/base/defaults/multiview_attention.py @@ -538,6 +538,21 @@ class MultiviewAttentionConfig: default="other_views", validator=attrs.validators.in_(("all_views", "own_view", "other_views")) ) + # Let a camera query's cross-view pass also key the other views' camera-pose action + # tokens, beside their vision tokens, under the same capture-time rule. Action tokens stay + # keys only: they never query the cross-view pass, and their own view's actions remain in + # the same-view pass alone. Requires maskless decomposed attention, no cross-view band + # expansion, no same-view temporal window and baseline geometry. False preserves existing + # recipes. + maskless_cross_view_attends_action: bool = False + + # Let camera-pose action tokens query the cross-view pass too, keyed only by the other + # views' action tokens under the same capture-time rule; they still key no vision token + # there, and their own view's actions stay in the same-view pass alone. Independent of + # ``maskless_cross_view_attends_action``, with the same requirements. False preserves + # existing recipes. + maskless_cross_view_action_attends_action: bool = False + # What the multiview backend lets the noisy tokens attend to. mask: MultiviewAttentionMaskConfig = attrs.field(factory=MultiviewAttentionMaskConfig) diff --git a/cosmos_framework/configs/base/defaults/reasoner.py b/cosmos_framework/configs/base/defaults/reasoner.py index 516c5ee22..b048cb712 100644 --- a/cosmos_framework/configs/base/defaults/reasoner.py +++ b/cosmos_framework/configs/base/defaults/reasoner.py @@ -183,7 +183,7 @@ class VLMConfig: class SoundUnderstandingConfig: """Optional audio-understanding pathway for standalone Reasoner/VLM models. - The encoder artifact initializes the frozen modality encoder independently + The encoder artifact initializes the optional trainable modality encoder independently of the Reasoner checkpoint. The audio projector is trainable by default and may be frozen explicitly after alignment. """ @@ -207,6 +207,9 @@ class SoundUnderstandingConfig: # ``separate_no_timestamps`` omits only audio-side timestamps, while # ``interleaved_av`` inserts each audio segment after its native video chunk. audio_layout: str = "separate_with_timestamps" + # Standalone audio remains enabled independently of video soundtrack extraction. + extract_audio_from_video: bool = False + freeze_encoder: bool = True freeze_projector: bool = False # Keep the full input embedding table as an optimizer parameter while # allowing gradients only through the audio start/end token positions. @@ -247,6 +250,8 @@ def validate_sound_understanding_config(config: Any, *, sound_und: bool) -> None if config.audio_layout not in _AUDIO_LAYOUTS: raise ValueError(f"sound_und_config.audio_layout must be one of {_AUDIO_LAYOUTS}, got {config.audio_layout!r}") for name, value in ( + ("extract_audio_from_video", config.extract_audio_from_video), + ("freeze_encoder", config.freeze_encoder), ("freeze_projector", config.freeze_projector), ("train_boundary_token_embeddings_only", config.train_boundary_token_embeddings_only), ("encoder_checkpoint_enabled", config.encoder_checkpoint_enabled), @@ -260,6 +265,8 @@ def validate_sound_understanding_config(config: Any, *, sound_und: bool) -> None raise ValueError("An enabled audio encoder checkpoint requires a non-empty path") if config.exclude_frozen_encoder_from_training_checkpoint and not config.encoder_checkpoint_enabled: raise ValueError("Excluding the frozen audio encoder from training checkpoints requires checkpoint loading") + if config.exclude_frozen_encoder_from_training_checkpoint and not config.freeze_encoder: + raise ValueError("A trainable audio encoder must be included in training checkpoints") # Configs for LLM models diff --git a/cosmos_framework/configs/base/reasoner/defaults/augmentors.py b/cosmos_framework/configs/base/reasoner/defaults/augmentors.py index 4a8872402..b776d6e37 100644 --- a/cosmos_framework/configs/base/reasoner/defaults/augmentors.py +++ b/cosmos_framework/configs/base/reasoner/defaults/augmentors.py @@ -6,6 +6,8 @@ from typing import Any +from omegaconf import OmegaConf + from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.data.generator.augmentors.reasoner.bytes_to_media import BytesToMedia from cosmos_framework.data.generator.augmentors.reasoner.filter_output_key import FilterOutputKey @@ -28,7 +30,15 @@ from cosmos_framework.utils.generator.deterministic_rng import DeterministicAugmentor +def _extract_video_audio(sound_und: bool, extract_audio_from_video: bool) -> bool: + for name, value in (("sound_und", sound_und), ("extract_audio_from_video", extract_audio_from_video)): + if not isinstance(value, bool): + raise TypeError(f"{name} must be a bool, got {type(value).__name__}") + return sound_und and extract_audio_from_video + + def create_data_augmentor_config() -> dict[str, Any]: + OmegaConf.register_new_resolver("reasoner_video_audio", _extract_video_audio, replace=True) config = { "bytes_to_media": L(BytesToMedia)( input_key="media", @@ -40,13 +50,17 @@ def create_data_augmentor_config() -> dict[str, Any]: video_timestamp_mode="${data_setting.video_timestamp_mode}", max_video_token_length="${data_setting.qwen_max_video_token_length}", # type: ignore processor=processor, - extract_audio="${model.config.sound_und}", + extract_audio=( + "${reasoner_video_audio:${model.config.sound_und}," + "${model.config.sound_und_config.extract_audio_from_video}}" + ), is_input_pickle_byptes=False, # If True, it means the input "media" is pickled bytes that needs to be unpickled first; if False, it means the input "media" is raw bytes that can be directly decoded to image/video. Set to False for most cases, and only set to True for some special datasets where media is stored as pickled bytes. ), # takes "videos" and output "videos" "prompt_format": L(PromptFormat)( # takes text_keys and output "conversation" input_keys=["texts"], text_chat_order="${data_setting.text_chat_order}", strip_thinking_prob="${data_setting.strip_thinking_prob}", + mask_inserted_reasoning_prefix="${data_setting.mask_inserted_reasoning_prefix}", ), "shuffle_text_media_order": L(ShuffleTextMediaOrder)(), "format_hot_fixes": L(FormatHotFixes)(), @@ -134,6 +148,7 @@ def create_data_augmentor_config() -> dict[str, Any]: data_stats_writer_enabled="${data_setting.data_stats_writer_enabled}", video_temporal_mode="${data_setting.qwen_video_temporal_mode}", video_timestamp_mode="${data_setting.video_timestamp_mode}", + keep_raw_media="${data_setting.keep_raw_media}", text_only=False, sound_und="${model.config.sound_und}", audio_encoder_type="${model.config.sound_und_config.audio_encoder_type}", diff --git a/cosmos_framework/configs/base/reasoner/defaults/callbacks.py b/cosmos_framework/configs/base/reasoner/defaults/callbacks.py index e6ad2f213..d0dde1fec 100644 --- a/cosmos_framework/configs/base/reasoner/defaults/callbacks.py +++ b/cosmos_framework/configs/base/reasoner/defaults/callbacks.py @@ -7,25 +7,28 @@ from hydra.core.config_store import ConfigStore +from cosmos_framework.callbacks.manual_gc import ManualGarbageCollection +from cosmos_framework.utils.lazy_config import PLACEHOLDER +from cosmos_framework.utils.lazy_config import LazyCall as L +from cosmos_framework.utils.callback import LowPrecisionCallback, WandBCallback from cosmos_framework.callbacks.data_stats_writer import DataStatsWriterCallback from cosmos_framework.callbacks.dataloader_state import DataLoaderStateCallback + from cosmos_framework.callbacks.grad_clip import GradClip from cosmos_framework.callbacks.hf_export import HFExportCallback from cosmos_framework.callbacks.iter_speed import IterSpeed from cosmos_framework.callbacks.learning_rate_logger import LearningRateLogger from cosmos_framework.callbacks.log_tensor_shape import LogTensorShapeCallback + from cosmos_framework.callbacks.loss_spike_rollback import LossSpikeRollback -from cosmos_framework.callbacks.manual_gc import ManualGarbageCollection from cosmos_framework.callbacks.param_count import ParamCount from cosmos_framework.callbacks.sampled_media_recorder import SampledMediaRecorder from cosmos_framework.callbacks.tokens_per_sec import VLMTokensPerSec from cosmos_framework.callbacks.wandb_log import WandbCallback as WandBCallbackMultiplier from cosmos_framework.callbacks.wandb_vis import VisualizationLoggingCallback + from cosmos_framework.callbacks.workflow_status import WorkflowStatusCallback from cosmos_framework.configs.base.defaults.job_monitor import JOB_MONITOR_CALLBACKS -from cosmos_framework.utils.callback import LowPrecisionCallback, WandBCallback -from cosmos_framework.utils.lazy_config import PLACEHOLDER -from cosmos_framework.utils.lazy_config import LazyCall as L # from cosmos_framework.utils.callback import NVTXCallback diff --git a/cosmos_framework/configs/base/reasoner/defaults/config.py b/cosmos_framework/configs/base/reasoner/defaults/config.py index d4ccd00c5..6be867469 100644 --- a/cosmos_framework/configs/base/reasoner/defaults/config.py +++ b/cosmos_framework/configs/base/reasoner/defaults/config.py @@ -14,10 +14,20 @@ class DataSetting: """Configuration for data. Attributes: + qwen_min_video_token_length: Minimum video token length per video. qwen_max_video_token_length: Maximum video token length. + qwen_max_video_frames: Optional deterministic decoded-frame ceiling per video. + qwen_adaptive_video_token_budget: Shrink video tokens when a row exceeds max_tokens. + qwen_adaptive_video_token_safety_margin: Tokens reserved after adaptive retokenization. + qwen_max_fps_thres: Maximum accepted native video FPS. qwen_target_fps: Target fps for video sampling. + keep_raw_media: Emit the debug-only raw_image / raw_video tensors; disable for long-context runs. text_chat_order: Order of text items in user messages. strip_thinking_prob: Per-sample probability of converting thinking data into non-thinking data. + offline_pack_text_media_shuffle_ratio: Offline packing only. Probability of swapping a user turn's + media and text, pinned to sample epoch 0 so measurement, restarts and training see the same order. + 0.01 matches the ordinary loader; 0 reproduces packs built without the swap. Part of the packing + contract: changing it needs a new packing plan (the Lance cache is reused). custom_system_prompt: System prompt injected when a conversation has no leading system message. strip_original_system_prompt: Remove existing system messages before optional custom prompt injection. distributor_type: "with_replace" (WeightedShardlistBasic) or "no_replace" (NoReplaceShardlistBasic). @@ -28,13 +38,19 @@ class DataSetting: long_threshold: Single seed-sample length threshold that emits the sample alone. data_stats_writer_enabled: Write per-sample parquet rows at every checkpoint save; a checkpoint counts as complete only once its rows are flushed and merged. - lance_cache_manifest: Recipe cache manifest for the shared remote Lance loader; None uses pinned PAI-v9. + lance_cache_manifest: Recipe cache manifest reusing pinned local or remote Lance sources. lance_data_root: Absolute local root of a receipt-backed conversation Lance export. """ + qwen_min_video_token_length: int = 16 qwen_max_video_token_length: int = 8192 + qwen_max_video_frames: int | None = None qwen_max_image_token_length: int = 8192 + qwen_adaptive_video_token_budget: bool = False + qwen_adaptive_video_token_safety_margin: int = 128 + qwen_max_fps_thres: float = 60.0 qwen_target_fps: float = 4.0 + keep_raw_media: bool = True use_native_edge_processor: bool = False video_timestamp_mode: str = attrs.field( default="qwen_index", @@ -56,6 +72,16 @@ class DataSetting: attrs.validators.le(1.0), ), ) + offline_pack_text_media_shuffle_ratio: float = attrs.field( + default=0.01, + validator=attrs.validators.and_( + attrs.validators.instance_of((int, float)), + attrs.validators.ge(0.0), + attrs.validators.le(1.0), + ), + ) + # Preserve processor-specific supervision unless reproducing an older packed run. + mask_inserted_reasoning_prefix: bool = False custom_system_prompt: str | None = "You are a helpful assistant." strip_original_system_prompt: bool = False temporal_localization_output_format: str = attrs.field( @@ -82,6 +108,55 @@ class DataSetting: data_stats_writer_enabled: bool = False lance_cache_manifest: str | None = None lance_data_root: str | None = None + # Safe opt-in switch used by ``lance_recipe_selectable``. The legacy explicit + # ``lance_offline_pack_recipe`` config remains available for exact resume compatibility. + offline_pack_enabled: bool = False + # Complete planned packs per rank/forward; distinct from the unpacked sample cap. + offline_pack_batch_size: int = attrs.field(default=1, validator=attrs.validators.ge(1)) + # Inert unless an offline-pack data config is selected. + offline_pack_recipe_manifest: str | None = None + offline_pack_schedule_path: str | None = None + offline_pack_validation_schedule_path: str | None = None + # Used only by the opt-in mixed-pack loader groups; existing defaults are unchanged. + offline_pack_mixture_schedule_paths: list[str] | None = None + offline_pack_recipe_source_manifest: str | None = None + offline_pack_recipe_output_root: str | None = None + # Opt-in persistent publication root; catalogs must already have been published offline. + offline_pack_artifact_root: str | None = None + # Optional integrity pin for the local recipe manifest. + # Required with an explicit published manifest; pin this alongside the URI on resume. + offline_pack_recipe_manifest_sha256: str | None = None + offline_pack_recipe_auto_prepare: bool = False + # Per-epoch launch requirement; does not change a published plan's content identity. + offline_pack_minimum_required_packs: int = attrs.field(default=0, validator=attrs.validators.ge(0)) + # The packed training resolver rejects setting both draw options before artifact I/O. + # Coverage computes total draws from all catalog rows plus weighted extras. + offline_pack_recipe_total_draws: int | None = None + offline_pack_recipe_weighted_extra_draws: int | None = attrs.field( + default=None, + validator=attrs.validators.optional(attrs.validators.ge(0)), + ) + offline_pack_recipe_required_training_epochs: int = attrs.field( + default=1, + validator=attrs.validators.ge(1), + ) + offline_pack_recipe_max_members_per_shard: int = attrs.field( + default=262_144, + validator=attrs.validators.ge(1), + ) + offline_pack_recipe_packing_workers: int = attrs.field( + default=1, + validator=attrs.validators.ge(1), + ) + offline_pack_recipe_global_ordering_mode: str = attrs.field( + default="epoch_permuted", + validator=attrs.validators.in_({"epoch_permuted", "legacy_direct"}), + ) + offline_pack_skip_recoverable_drops: bool = False + offline_pack_max_consecutive_skipped_packs: int = attrs.field( + default=32, + validator=attrs.validators.ge(1), + ) @attrs.define(slots=False) diff --git a/cosmos_framework/configs/base/reasoner/defaults/policy_config.py b/cosmos_framework/configs/base/reasoner/defaults/policy_config.py index 83191332a..bde6427a2 100644 --- a/cosmos_framework/configs/base/reasoner/defaults/policy_config.py +++ b/cosmos_framework/configs/base/reasoner/defaults/policy_config.py @@ -50,9 +50,22 @@ class PolicyConfig: # must use the explicit fields above so PEFT equivalence can be validated. lora: Union[str, None] = None enable_liger_kernel: bool = False - # Dense Qwen3.5 caption opt-ins; other recipes retain the standard logits loss. + # Fused loss supports dense Qwen3.5 and dense/MoE Qwen3-VL; default recipes retain logits loss. enable_fused_weighted_ce: bool = False + # Recompute scalar CE during backward instead of retaining one dense head gradient per conversation. + recompute_fused_weighted_ce: bool = False + # Packed Qwen3-VL: sum native per-conversation head gradients in FP32, then cast once. + qwen3_vl_accumulate_head_grad_fp32: bool = True + # The recurrent precision option is specific to dense Qwen3.5 caption models. qwen35_fp32_recurrent_a_log: bool = False + # Experimental numerical stabilization; preserves expert parameter/checkpoint layouts. + qwen3_vl_moe_fixed_sum: bool = False + # Experimental CP: replicate vision encoding, then shard merged decoder inputs. + qwen3_vl_cp_multimodal: bool = False + # Bound vision varlen backward workspaces; zero preserves one attention call. + qwen3_vl_vision_attention_max_padded_tokens: int = 0 + # Group complete packed conversations to bound decoder varlen backward workspaces. + qwen3_vl_decoder_attention_max_padded_tokens: int = 0 trainable_map: Union[str, None] = None monkey_patch_for_text_only_data: bool = False diff --git a/cosmos_framework/configs/simulation/config.py b/cosmos_framework/configs/simulation/config.py index efc3a5b16..276c8bb05 100644 --- a/cosmos_framework/configs/simulation/config.py +++ b/cosmos_framework/configs/simulation/config.py @@ -14,22 +14,16 @@ from hydra.core.config_store import ConfigStore +from cosmos_framework.utils.lazy_config import LazyCall as L +from cosmos_framework.utils.lazy_config import LazyDict +from cosmos_framework.trainer import ImaginaireTrainer +from cosmos_framework.utils.config_helper import import_all_modules_from_package, import_module from cosmos_framework.configs.base.config import Config as VFMConfig from cosmos_framework.configs.base.defaults.callbacks import register_callbacks from cosmos_framework.configs.base.defaults.checkpointer import ( register_checkpoint, register_ckpt_type, ) -from cosmos_framework.configs.base.defaults.distill_callbacks import ( - register_distill_callbacks, -) -from cosmos_framework.configs.base.defaults.distill_ckpt_type import ( - register_distill_ckpt_type, -) -from cosmos_framework.configs.base.defaults.distill_optimizer import ( - register_optimizer, - register_scheduler, -) from cosmos_framework.configs.base.defaults.ema import register_ema from cosmos_framework.configs.base.defaults.model import register_model from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig @@ -40,6 +34,16 @@ register_sound_tokenizer, register_tokenizer, ) +from cosmos_framework.configs.base.defaults.distill_callbacks import ( + register_distill_callbacks, +) +from cosmos_framework.configs.base.defaults.distill_ckpt_type import ( + register_distill_ckpt_type, +) +from cosmos_framework.configs.base.defaults.distill_optimizer import ( + register_optimizer, + register_scheduler, +) from cosmos_framework.model.generator.consistency_model.tf_dcm import ( TeacherForcingDCMModel, TeacherForcingDCMModelConfig, @@ -48,17 +52,13 @@ SelfForcingDMDConfig, SelfForcingDMDModel, ) +from cosmos_framework.trainer.distillation import ( + DistillationTrainer, # noqa: F401 - release-only binding of the distillation trainer +) from cosmos_framework.model.generator.omni_mot_causal_model import ( OmniMoTCausalModel, OmniMoTCausalModelConfig, ) -from cosmos_framework.trainer import ImaginaireTrainer -from cosmos_framework.trainer.distillation import ( - DistillationTrainer, # noqa: F401 - release-only binding of the distillation trainer -) -from cosmos_framework.utils.config_helper import import_all_modules_from_package, import_module -from cosmos_framework.utils.lazy_config import LazyCall as L -from cosmos_framework.utils.lazy_config import LazyDict # Experiment registries imported when make_config() is called without an explicit # registry. The released copy imports only the families that ship with it. diff --git a/cosmos_framework/data/generator/action/action_caption_attribute_adapter.py b/cosmos_framework/data/generator/action/action_caption_attribute_adapter.py index 1b761c9e2..187116568 100644 --- a/cosmos_framework/data/generator/action/action_caption_attribute_adapter.py +++ b/cosmos_framework/data/generator/action/action_caption_attribute_adapter.py @@ -26,6 +26,7 @@ } ViewComposition = Literal[ + "separate_camera_views", "front", "above", "third_person", @@ -40,6 +41,7 @@ "above_over_side_third_person", ] ViewCompositionMappings: dict[ViewComposition, str] = { + "separate_camera_views": "The synchronized camera views are encoded separately, in a fixed camera order.", "front": "This video is captured from a dynamic front-facing perspective looking at the scene.", "above": "This video is captured from a third-person perspective looking towards the agent from above.", "third_person": "This video is captured from a third-person perspective looking towards the robot.", @@ -56,6 +58,9 @@ CaptionSubject = Literal[ "camera_ego_pose", + "right_wrist_camera", + "dual_wrist_cameras", + "ego_and_dual_wrist_cameras", "vehicle_ego_pose", "right_arm_end_effector", "right_arm_attachment_site_end_effector", @@ -69,6 +74,18 @@ "human_head_wrists_and_fingertips", ] CaptionSubjectMappings: dict[CaptionSubject, str] = { + "right_wrist_camera": ( + "The action target is the wrist-mounted camera motion and gripper openness, " + "mapped to the right wrist-pose and right-openness components." + ), + "dual_wrist_cameras": ( + "The action targets are the right and left wrist-mounted camera motions and gripper openness, " + "mapped to the right and left wrist-pose and openness components." + ), + "ego_and_dual_wrist_cameras": ( + "The action targets are the head ego motion, right and left wrist-mounted camera motions, and " + "gripper openness, mapped to the ego-pose, right and left wrist-pose and openness components." + ), "dual_arm_joint_positions": ( "The action performed by the robot is defined as the joint positions of its left and right arms " "and both grippers' positions." diff --git a/cosmos_framework/data/generator/action/utils/camera_views.py b/cosmos_framework/data/generator/action/utils/camera_views.py new file mode 100644 index 000000000..9a58af968 --- /dev/null +++ b/cosmos_framework/data/generator/action/utils/camera_views.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Physical camera names shared by native geometry and per-view Action packing.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from cosmos_framework.data.generator.multiview.camera_pose_actions import CameraActionSlot + +# These are physical camera names from the native Lance adapters, not inferred +# left/right substrings (DROID's left wrist lens is on its single right arm). +CAMERA_VIEW_ACTION_SLOTS: dict[str, CameraActionSlot] = { + "observation.image.wrist_image_left": "right_wrist", + "observation.image.exterior_image_1_left": "ego", + "observation.image.exterior_image_2_left": "ego", + "observation.images.top_head": "ego", + "observation.images.head": "ego", + "observation.images.hand_left": "left_wrist", + "observation.images.hand_right": "right_wrist", + "observation.images.top": "ego", + "observation.images.left": "left_wrist", + "observation.images.right": "right_wrist", + "observation.images.top_camera-images-rgb": "ego", + "observation.images.left_camera-images-rgb": "left_wrist", + "observation.images.right_camera-images-rgb": "right_wrist", + "observation.images.left_wrist": "left_wrist", + "observation.images.right_wrist": "right_wrist", + "observation.images.faceImg": "ego", + "observation.images.leftImg": "left_wrist", + "observation.images.rightImg": "right_wrist", + "observation.image.left_main_camera_rgb": "left_wrist", + "observation.image.right_main_camera_rgb": "right_wrist", + "observation.rgb.zed_link_camera_0": "ego", + "observation.rgb.left_realsense_link_camera_0": "left_wrist", + "observation.rgb.right_realsense_link_camera_0": "right_wrist", + "observation.images.cam_high": "ego", + "observation.images.cam_left_wrist": "left_wrist", + "observation.images.cam_right_wrist": "right_wrist", + "observation.images.camera_top_left": "ego", + "observation.images.camera_wrist_left": "left_wrist", + "observation.images.camera_wrist_right": "right_wrist", +} +_CAMERA_KEY_BY_SLOT: dict[CameraActionSlot, str] = { + "ego": "camera_head", + "left_wrist": "camera_wrist_left", + "right_wrist": "camera_wrist_right", +} + +# Physical roles are separate from action slots: exterior, head and fixed top +# cameras can all occupy ego slots without being the same kind of camera. +CAMERA_CAPTION_PREFIXES: dict[str, str] = { + "camera_wrist_right": "The view is from the camera mounted on the right wrist.", + "camera_wrist_left": "The view is from the camera mounted on the left wrist.", + "camera_head": "The view is from the camera mounted on the robot head.", + "camera_exterior_1": "The view is from the first static exterior camera.", + "camera_exterior_2": "The view is from the second static exterior camera.", + "camera_top": "The view is from the fixed camera above the workstation.", + "camera_front_static": "The view is from the fixed front camera.", +} +_PHYSICAL_CAMERA_KEY_BY_NATIVE_VIEW: dict[str, str] = { + view: _CAMERA_KEY_BY_SLOT[slot] for view, slot in CAMERA_VIEW_ACTION_SLOTS.items() if slot != "ego" +} +_PHYSICAL_CAMERA_KEY_BY_NATIVE_VIEW.update( + { + "observation.image.exterior_image_1_left": "camera_exterior_1", + "observation.image.exterior_image_2_left": "camera_exterior_2", + "observation.images.top_head": "camera_head", + "observation.images.head": "camera_head", + "observation.images.top": "camera_top", + "observation.images.top_camera-images-rgb": "camera_top", + "observation.images.faceImg": "camera_front_static", + "observation.rgb.zed_link_camera_0": "camera_head", + "observation.images.cam_high": "camera_top", + "observation.images.camera_top_left": "camera_head", + } +) + + +def native_camera_caption_prefixes(view_names: Sequence[str]) -> dict[str, str]: + """Describe known physical cameras; never guess a role for an unknown stream.""" + return { + view: CAMERA_CAPTION_PREFIXES[_PHYSICAL_CAMERA_KEY_BY_NATIVE_VIEW[view]] + for view in view_names + if view in _PHYSICAL_CAMERA_KEY_BY_NATIVE_VIEW + } + + +def camera_view_key(view_name: str) -> str: + """Resolve a native observation name to its canonical optical camera key.""" + return _CAMERA_KEY_BY_SLOT[CAMERA_VIEW_ACTION_SLOTS[view_name]] + + +def camera_action_caption_subject(view_name: str, slot: CameraActionSlot | None) -> str: + """Describe only the action components allocated to this physical view.""" + if slot is None: + return "There is no action target." + if slot == "ego": + subject = ( + "head ego motion" + if _PHYSICAL_CAMERA_KEY_BY_NATIVE_VIEW[view_name] == "camera_head" + else "camera ego motion" + ) + return f"The action targets are the {subject}, mapped to the ego-pose." + side = "right" if slot == "right_wrist" else "left" + return ( + f"The action targets are the {side} wrist-mounted camera motion and gripper openness, " + f"mapped to the {side} wrist-pose and {side}-openness components." + ) diff --git a/cosmos_framework/data/generator/action/utils/multiview.py b/cosmos_framework/data/generator/action/utils/multiview.py new file mode 100644 index 000000000..e027e4b48 --- /dev/null +++ b/cosmos_framework/data/generator/action/utils/multiview.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Bridge native Action windows to the multiview unified-wrist-slot layout. + +Camera poses have already been converted from end-effector targets at the +sampler boundary. Reuse !14785's slot scatter, metric-gripper normalization and +codec, retaining one camera-major row block per view and all native channels +for the single-view human streams. +""" + +from collections.abc import Mapping +from typing import Any, cast + +import numpy as np +import torch + +from cosmos_framework.data.generator.action.utils.camera_views import ( + CAMERA_VIEW_ACTION_SLOTS, + camera_action_caption_subject, +) +from cosmos_framework.data.generator.action.utils.pose_utils import PoseConvention +from cosmos_framework.data.generator.action.utils.unified_action_schema import ( + EGO_POSE, + LEFT_OPENNESS, + LEFT_WRIST_POSE, + RIGHT_OPENNESS, + RIGHT_WRIST_POSE, + UNIFIED_ACTION_DIM, +) +from cosmos_framework.data.generator.multiview.camera_pose_actions import ( + CameraActionSlot, + CameraActionSlotBinding, + CameraActionSlots, + build_unified_camera_slot_action, +) + +_POSE_SLOTS: dict[CameraActionSlot, slice] = { + "ego": EGO_POSE, + "right_wrist": RIGHT_WRIST_POSE, + "left_wrist": LEFT_WRIST_POSE, +} + + +def _camera_action_slot( + camera: str, + camera_pose_actions: Mapping[str, np.ndarray | None] | None, + valid_mask: torch.Tensor | None, # [D_raw] +) -> CameraActionSlot | None: + """Share the target-availability decision between captions and slot projection.""" + if camera not in CAMERA_VIEW_ACTION_SLOTS: + raise ValueError(f"No audited wrist/ego slot mapping for camera {camera!r}") + slot = CAMERA_VIEW_ACTION_SLOTS[camera] + if camera_pose_actions is not None: + return slot if camera_pose_actions[camera] is not None else None + if valid_mask is not None and not bool(valid_mask[_POSE_SLOTS[slot]].all()): + return None + return slot + + +def camera_action_caption_subjects( + sample: dict[str, Any], + action_valid_mask: torch.Tensor | None, # [D_raw] +) -> list[str]: + """Render camera-major subjects from the same bindings used by the action projection.""" + return [ + camera_action_caption_subject( + camera, _camera_action_slot(camera, sample.get("camera_pose_actions"), action_valid_mask) + ) + for camera in sample["camera_view_names"] + ] + + +def split_action_by_camera(sample: dict[str, Any]) -> dict[str, Any]: + """Select each camera's existing canonical slots without replicating both arms.""" + camera_names = sample.pop("camera_view_names") + pose_convention = sample.pop("camera_action_pose_convention") + gripper_width = sample.pop("camera_action_gripper_width_m") + camera_pose_actions = sample.pop("camera_pose_actions", None) + if len(camera_names) == 1: + # Native human, camera/AV and single-wrist targets already occupy their + # trained slots, including fingertips, route/history and grippers. + return sample + raw = sample["action_raw"] + if raw.shape[-1] != UNIFIED_ACTION_DIM: + raise ValueError("Per-camera wrist slots require the unified_v1 Action schema") + valid_mask = sample["action_valid_mask"][:UNIFIED_ACTION_DIM] + per_view_poses: list[np.ndarray | None] = [] + bindings: dict[str, CameraActionSlotBinding] = {} + missing_views: list[int] = [] + for view, camera in enumerate(camera_names): + slot = _camera_action_slot(camera, camera_pose_actions, valid_mask) + if slot is None: + per_view_poses.append(None) + missing_views.append(view) + continue + measured_pose = ( + camera_pose_actions[camera] if camera_pose_actions is not None else raw[:, _POSE_SLOTS[slot]].cpu().numpy() + ) # [T,9] + per_view_poses.append(measured_pose) + if slot == "ego": + bindings[camera] = CameraActionSlotBinding(slot="ego") + else: + openness_slot = RIGHT_OPENNESS if slot == "right_wrist" else LEFT_OPENNESS + if not bool(valid_mask[openness_slot]): + raise ValueError(f"Wrist camera {camera!r} has no calibrated gripper openness") + bindings[camera] = CameraActionSlotBinding( + slot=slot, + openness=raw[:, openness_slot].cpu().numpy(), + openness_frame_offset=0, + ) + projected = build_unified_camera_slot_action( + per_view_poses, + frame_indices=list(range(raw.shape[0] + 1)), + action_slots=CameraActionSlots(bindings, gripper_width), + camera_keys=camera_names, + pose_convention=cast(PoseConvention, pose_convention), + max_action_dim=sample["action"].shape[-1], + ) + sample.update( + action=projected.action, + action_raw=projected.action_raw, + raw_action_dim=projected.raw_action_dim, + action_valid_mask=projected.action_valid_mask, + action_processing_record=projected.action_processing_record, + num_views_per_action_item=torch.tensor(len(camera_names), dtype=torch.int64), + ) + plan = sample["sequence_plan"] + source_conditions = list(plan.condition_frame_indexes_action) + rows_per_view = raw.shape[0] + plan.condition_frame_indexes_action = [ + view * rows_per_view + frame + for view in range(len(camera_names)) + for frame in (range(rows_per_view) if view in missing_views else source_conditions) + ] + return sample diff --git a/cosmos_framework/data/generator/action/utils/transforms.py b/cosmos_framework/data/generator/action/utils/transforms.py index 08c1fafdb..0fd9e8ea6 100644 --- a/cosmos_framework/data/generator/action/utils/transforms.py +++ b/cosmos_framework/data/generator/action/utils/transforms.py @@ -32,16 +32,22 @@ ActionProcessor, ) from cosmos_framework.data.generator.action.utils.json_formatter import ActionPromptJsonFormatter +from cosmos_framework.data.generator.action.utils.multiview import ( + camera_action_caption_subjects, + split_action_by_camera, +) from cosmos_framework.data.generator.action.utils.viewpoint_utils import ViewpointTextInfo from cosmos_framework.data.generator.augmentors.duration_fps_text_timestamps import DurationFPSTextTimeStamps from cosmos_framework.data.generator.augmentors.idle_frames_text_info import IdleFramesTextInfo from cosmos_framework.data.generator.augmentors.resolution_text_info import ResolutionTextInfo from cosmos_framework.data.generator.augmentors.text_tokenizer import TextTokenizerTransform +from cosmos_framework.data.generator.multiview.caption_format import format_separate_view_captions from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO from cosmos_framework.data.generator.sequence_packing import SequencePlan from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution ACTION_MODE_DESCRIPTIONS: dict[str, str] = { + "policy": "Predict future action given the image.", "forward_dynamics": "Predict the future video given the image and action.", "inverse_dynamics": "Predict the underlying action given the video.", "wam": "Predict the future video and action given the image.", @@ -521,6 +527,11 @@ class ActionTransformPipeline: Requires an ``action_caption_attributes`` mapping on every sample; missing metadata raises before resizing or tokenization, including JSON prompts. Defaults to ``False``. + append_camera_role_to_caption: Repeat the formatted task prompt per + separate camera clip, prepend its physical role when known, and + tokenize each caption separately. With plain-caption semantics, + describe only that view's allocated action targets. Raw dataset + camera names are never added to the prompt. Defaults to ``False``. append_duration_fps_timestamps: Whether to append duration and FPS metadata to the caption (matching VFM's ``DurationFPSTextTimeStamps`` augmentor). Defaults to ``True``. @@ -575,6 +586,7 @@ def __init__( format_prompt_as_json: bool = False, enable_mode_specific_prompt: bool = False, format_prompt_float_seconds: bool = False, + append_camera_role_to_caption: bool = False, ) -> None: self.caption_key: str = caption_key self.video_temporal_downsample: int = video_temporal_downsample @@ -582,6 +594,7 @@ def __init__( self.action_channel_masking: bool = action_channel_masking self.enable_mode_specific_prompt: bool = enable_mode_specific_prompt self.append_action_caption_semantics: bool = append_action_caption_semantics + self.append_camera_role_to_caption: bool = append_camera_role_to_caption self.action_processor: ActionProcessor = ActionProcessor( max_action_dim=max_action_dim, action_channel_masking=action_channel_masking, @@ -658,10 +671,13 @@ def __init__( if tokenizer_config is not None: self.text_tokenizer = TextTokenizerTransform( input_keys=[caption_key], - output_keys=[text_token_key], + output_keys=[text_token_key, "text_token_lengths"] + if append_camera_role_to_caption + else [text_token_key], args={ "tokenizer_config": tokenizer_config, "cfg_dropout_rate": cfg_dropout_rate, + "tokenize_separately": append_camera_role_to_caption, }, ) @@ -697,6 +713,59 @@ def _apply_plain_caption_postfixes(self, data_dict: dict, mode: object) -> dict: return data_dict + def _prepare_video(self, data_dict: dict, resolution: str | None) -> tuple[dict, int | None]: + """Resize camera clips onto one canvas while preserving the full prompt horizon.""" + camera_videos = data_dict.pop("camera_videos", None) + if camera_videos is None: + return self.video_resize(data_dict, resolution), None + + if not camera_videos or any(view.ndim != 4 for view in camera_videos): + raise ValueError("camera_videos must contain nonempty [C,T,H,W] camera clips") + num_views = len(camera_videos) + frames_per_view = camera_videos[0].shape[1] + if any(view.shape[:2] != camera_videos[0].shape[:2] for view in camera_videos): + raise ValueError("All camera clips must have the same channels and synchronized frame count") + first = self.video_resize({"video": camera_videos[0]}, resolution) + target_h, target_w = first["video"].shape[-2:] + resized_views = [first["video"]] + for view in camera_videos[1:]: + resized = reflection_pad_to_target({"video": view}, ["video"], self.keep_aspect_ratio, target_w, target_h) + resized_views.append(resized["video"]) + data_dict["video"] = torch.cat(resized_views, dim=1).contiguous() # [C,V*T,H,W] + # One shared canvas contains each independently padded camera. Do not + # crop every camera according to the first camera's original aspect. + data_dict["image_size"] = torch.tensor([target_h, target_w, target_h, target_w], dtype=torch.float) + data_dict["video_num_frames"] = frames_per_view + return data_dict, num_views + + def _finalize_video_layout(self, data_dict: dict, mode: object, num_views: int | None) -> int: + """Select Policy observations after prompt formatting and attach the camera layout.""" + video = data_dict.get("video") + assert video is not None, "video is required" + video_length = int(data_dict.pop("video_num_frames", video.shape[1])) # logical T before decode compaction + if mode == "policy": + # Prompt formatting already describes the full action horizon. The model + # conditions only on the observation image, with no future RGB targets. + if num_views is None: + data_dict["video"] = video[:, :1].contiguous() + else: + data_dict["video"] = ( + video.reshape(video.shape[0], num_views, video_length, *video.shape[-2:])[:, :, :1] + .reshape(video.shape[0], num_views, *video.shape[-2:]) + .contiguous() + ) + video_length = 1 + if num_views is not None: + data_dict["sample_n_views"] = torch.tensor(num_views, dtype=torch.int64) + data_dict["num_video_frames_per_view"] = torch.tensor(video_length, dtype=torch.int64) + data_dict["enable_per_camera_vae_encoding"] = True + data_dict["view_indices_selection"] = torch.arange(num_views, dtype=torch.int64) + data_dict["view_indices"] = torch.arange(num_views, dtype=torch.int64).repeat_interleave(video_length) + # Single-view targets keep their native layout. The wrist-slot + # projection below replaces this count for multi-camera samples. + data_dict["num_views_per_action_item"] = torch.tensor(1, dtype=torch.int64) + return video_length + def __call__( self, data_dict: dict, @@ -728,7 +797,8 @@ def __call__( 7. Append idle-frame metadata (Pi0.7-style) to plain prompts unless the sample is in inverse dynamics mode (if enabled). 8. Tokenize caption text (if enabled). - 9. Build a ``SequencePlan`` from the ``"mode"`` key (if present). + 9. Keep only the observation image for Policy, retaining the action + horizon and prompt duration, and build a ``SequencePlan``. 10. Preserve the canonical unnormalized and unpadded action as ``"action_raw"``, normalize native actions or consume an explicitly encoded normalized action, pad ``"action"`` to ``max_action_dim``, and @@ -767,7 +837,7 @@ def __call__( raise TypeError("action_caption_attributes must be a mapping.") # 1. Resize + reflection-pad spatial dimensions to the closest predefined target from ``VIDEO_RES_SIZE_INFO[resolution]``. - data_dict = self.video_resize(data_dict, resolution) + data_dict, num_views = self._prepare_video(data_dict, resolution) # 2. Format the caption as structured JSON when requested; otherwise run the legacy string appenders. if self.prompt_json_formatter is not None: @@ -787,21 +857,46 @@ def __call__( data_dict[self.caption_key] = add_action_mode_metadata(data_dict.get(self.caption_key), mode) # 4-7. Append plain-caption metadata in canonical order. - data_dict = self._apply_plain_caption_postfixes(data_dict, mode) + if self.append_camera_role_to_caption and self.append_action_caption_semantics and (num_views or 0) > 1: + attributes = data_dict["action_caption_attributes"] + captions = [] + for subject in camera_action_caption_subjects(data_dict, action_valid_mask): + view_data = { + **data_dict, + "action_caption_attributes": {**attributes, "subject_postfix": subject}, + } + captions.append(self._apply_plain_caption_postfixes(view_data, mode)[self.caption_key]) + data_dict[self.caption_key] = captions + else: + data_dict = self._apply_plain_caption_postfixes(data_dict, mode) # This nested dictionary cannot be consumed by the training collator. # Its semantic fields have already been rendered into the caption. data_dict.pop("action_caption_attributes", None) # 8. Tokenize caption text (if enabled). + camera_prefixes = data_dict.pop("camera_caption_prefixes", None) + if self.append_camera_role_to_caption: + camera_names = data_dict.get("camera_view_names") + if num_views is None or not camera_names or len(camera_names) != num_views: + raise ValueError("Per-view captions require separate camera clips and matching camera_view_names") + # Native Action shares the task caption, with each camera's own + # action subject, preserving the camera-major video/action order. + captions = data_dict[self.caption_key] + data_dict[self.caption_key] = format_separate_view_captions( + captions if isinstance(captions, list) else [captions] * num_views, + camera_names=camera_names, + append_camera_role_to_caption=True, + camera_prefixes=camera_prefixes, + ) + # The per-view tokenizer validates this before video finalization. + data_dict["sample_n_views"] = torch.tensor(num_views, dtype=torch.int64) if self.text_tokenizer is not None: data_dict = self.text_tokenizer(data_dict) # 9. Build a ``SequencePlan`` from the ``"mode"`` key (if present). - video = data_dict.get("video") + video_length = self._finalize_video_layout(data_dict, mode, num_views) action = data_dict.get("action") - assert video is not None, "video is required" - video_length = int(data_dict.pop("video_num_frames", video.shape[1])) # logical T before decode compaction action_length = action.shape[0] if isinstance(action, torch.Tensor) else max(video_length - 1, 0) # Prepend history action frames (ground-truth conditioning) if present. @@ -842,5 +937,7 @@ def __call__( normalized_action=normalized_action, action_codec=action_codec, ) + if num_views is not None: + data_dict = split_action_by_camera(data_dict) return data_dict diff --git a/cosmos_framework/data/generator/action/utils/transforms_test.py b/cosmos_framework/data/generator/action/utils/transforms_test.py index f1466b9bf..bd6b0aa5b 100644 --- a/cosmos_framework/data/generator/action/utils/transforms_test.py +++ b/cosmos_framework/data/generator/action/utils/transforms_test.py @@ -3,11 +3,18 @@ from __future__ import annotations +import copy from types import SimpleNamespace +from typing import Any import pytest import torch +from cosmos_framework.data.generator.action.action_caption_attribute_adapter import ( + CaptionSubjectMappings, + ViewCompositionMappings, +) +from cosmos_framework.data.generator.action.utils.camera_views import native_camera_caption_prefixes from cosmos_framework.data.generator.action.utils.json_formatter import ActionPromptJsonFormatter from cosmos_framework.data.generator.action.utils.transforms import ( ActionTransformPipeline, @@ -15,8 +22,212 @@ reflection_pad_to_target, remove_reflection_padding, ) +from cosmos_framework.data.generator.action.utils.unified_action_schema import UNIFIED_ACTION_DIM from cosmos_framework.data.generator.augmentors.duration_fps_text_timestamps import DurationFPSTextTimeStamps from cosmos_framework.data.generator.augmentors.resolution_text_info import ResolutionTextInfo +from cosmos_framework.data.generator.joint_dataloader import custom_collate_fn +from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean +from cosmos_framework.data.generator.sequence_packing.packers import pack_input_sequence + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("mode", ["policy", "wam", "inverse_dynamics"]) +@pytest.mark.parametrize("num_views", [1, 2]) +@pytest.mark.parametrize("json_prompt", [False, True]) +@pytest.mark.parametrize("cfg_dropped", [False, True]) +def test_action_camera_roles_preserve_camera_layout_and_task_prompt( + monkeypatch: pytest.MonkeyPatch, mode: str, num_views: int, json_prompt: bool, cfg_dropped: bool +) -> None: + class Processor: + def tokenize_text(self, caption: str, **kwargs: object) -> list[int]: + return list(caption.encode("utf-8")) or [0] + + monkeypatch.setattr( + "cosmos_framework.data.generator.augmentors.text_tokenizer.lazy_instantiate", lambda _: Processor() + ) + names = ("observation.images.right_wrist", "observation.images.left_wrist")[:num_views] + videos = [torch.full((3, 17, 4, 6), float(index)) for index in range(num_views)] + action = torch.zeros(16, UNIFIED_ACTION_DIM) + # Identity rotations in each canonical pose slot; openness is valid in both wrists. + action[:, [3, 7, 12, 16, 37, 41]] = 1.0 + action[:, [33, 58]] = 0.5 + data = { + "ai_caption": "Move the camera.", + "video": torch.cat(videos, dim=1), + "camera_videos": videos, + "camera_view_names": names, + "camera_caption_prefixes": native_camera_caption_prefixes(names), + "camera_action_pose_convention": "backward_chunk_anchored_16f", + "camera_action_gripper_width_m": 1.0, + "action": action, + "conditioning_fps": torch.tensor(8), + "mode": mode, + "viewpoint": "separate_camera_views", + } + options: dict[str, Any] = dict( + pad_keys=[], + max_action_dim=64, + append_viewpoint_info=True, + append_resolution_info=False, + format_prompt_as_json=json_prompt, + enable_mode_specific_prompt=True, + ) + valid_mask = torch.ones(UNIFIED_ACTION_DIM, dtype=torch.bool) + baseline = ActionTransformPipeline(**options)(copy.deepcopy(data), "256", action_valid_mask=valid_mask) + pipeline = ActionTransformPipeline( + **options, tokenizer_config={}, cfg_dropout_rate=0.1, append_camera_role_to_caption=True + ) + monkeypatch.setattr("random.random", lambda: 0.05 if cfg_dropped else 0.9) + result = pipeline(data, "256", action_valid_mask=valid_mask) + + if cfg_dropped: + assert result["ai_caption"] == [""] * num_views + else: + for name, caption in zip(names, result["ai_caption"]): + assert caption.startswith(f"{native_camera_caption_prefixes(names)[name]}\n\n") + assert "View:" not in caption + assert all(camera_name not in caption for camera_name in names) + assert ("Move the camera." in caption) is (mode != "inverse_dynamics") + assert "2s" in caption if json_prompt else "2.0 seconds" in caption + assert "separately" in caption + assert "concatenated" not in caption + assert result["text_token_lengths"] == [len(tokens) for tokens in result["text_token_ids"]] + assert len(result["text_token_ids"]) == int(result["sample_n_views"]) == num_views + torch.testing.assert_close(result["video"], baseline["video"]) + torch.testing.assert_close(result["action"], baseline["action"]) + assert result["sequence_plan"] == baseline["sequence_plan"] + batch = custom_collate_fn([result]) + assert "camera_caption_prefixes" not in result + assert batch["ai_caption"] == [result["ai_caption"]] + assert batch["text_token_lengths"] == [result["text_token_lengths"]] + # The model assigns local view IDs from these ragged caption boundaries. + # Exercise the actual packer with that contract and camera-major VAE/action tensors. + plan = copy.deepcopy(result["sequence_plan"]) + plan.text_view_ids = list(range(len(batch["text_token_lengths"][0]))) + latent_frames = 1 if mode == "policy" else 5 + clean = GenerationDataClean( + batch_size=1, + is_image_batch=False, + x0_tokens_vision=[torch.zeros(1, 4, num_views * latent_frames, 1, 1)], + num_vision_items_per_sample=[1], + num_views_per_vision_item=[num_views], + x0_tokens_action=[result["action"]], + num_views_per_action_item=[int(result["num_views_per_action_item"])], + fps_vision=torch.tensor([8.0]), + fps_action=torch.tensor([8.0]), + ) + packed = pack_input_sequence( + [plan], + [tokens.tolist() for tokens in result["text_token_ids"]], + clean, + torch.tensor([0.5]), + {"eos_token_id": 256, "start_of_generation": 257, "end_of_generation": 258}, + action_dim=64, + ) + assert packed.text_caption_view_ids == [list(range(num_views))] + assert len(packed.text_caption_lens[0]) == num_views + assert packed.num_views_per_vision_item == packed.num_views_per_action_item == [num_views] + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_action_camera_roles_reject_missing_separate_camera_metadata() -> None: + pipeline = ActionTransformPipeline( + pad_keys=[], + append_viewpoint_info=False, + append_duration_fps_timestamps=False, + append_resolution_info=False, + append_camera_role_to_caption=True, + ) + with pytest.raises(ValueError, match="require separate camera clips"): + pipeline({"ai_caption": "move", "video": torch.zeros(3, 17, 4, 6), "mode": "policy"}, "256") + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("mode", ["policy", "wam", "forward_dynamics", "inverse_dynamics"]) +@pytest.mark.parametrize("head_mapped", [False, True]) +@pytest.mark.parametrize("explicit_camera_poses", [False, True]) +def test_per_view_action_subjects_match_projected_targets( + mode: str, head_mapped: bool, explicit_camera_poses: bool +) -> None: + # Exercise B1K's native head/left/right order, including an explicitly + # unavailable optical head pose overriding an otherwise valid legacy head. + names = ( + "observation.rgb.zed_link_camera_0", + "observation.rgb.left_realsense_link_camera_0", + "observation.rgb.right_realsense_link_camera_0", + ) + videos = [torch.zeros(3, 17, 4, 6) for _ in names] # 3 x [C,T,H,W] + action = torch.zeros(16, UNIFIED_ACTION_DIM) # [T,D] + action[:, [3, 7, 12, 16, 37, 41]] = 1.0 # [T,6] identity rotations + action[:, [33, 58]] = 0.5 # [T,2] gripper openness + attributes = { + "domain_postfix": "The video is captured from a synthetic environment.", + "embodiment_postfix": "The video shows a dual-arm robot.", + "view_postfix": ViewCompositionMappings["separate_camera_views"], + "subject_postfix": CaptionSubjectMappings["ego_and_dual_wrist_cameras"], + } + data = { + "ai_caption": "The goal is clean dining table and clear crumbs.", + "video": torch.cat(videos, dim=1), # [C,3*T,H,W] + "camera_videos": videos, + "camera_view_names": names, + "camera_caption_prefixes": native_camera_caption_prefixes(names), + "camera_action_pose_convention": "backward_chunk_anchored_16f", + "camera_action_gripper_width_m": 0.1, + "action": action, + "action_caption_attributes": attributes, + "conditioning_fps": torch.tensor(8), # [] + "mode": mode, + "viewpoint": "separate_camera_views", + } + valid_mask = torch.ones(UNIFIED_ACTION_DIM, dtype=torch.bool) # [D] + if explicit_camera_poses: + data["camera_pose_actions"] = { + names[0]: action[:, :9].numpy() if head_mapped else None, # [T,9] or None + names[1]: action[:, 34:43].numpy(), # [T,9] + names[2]: action[:, 9:18].numpy(), # [T,9] + } + else: + valid_mask[:9] = head_mapped # [9] + pipeline = ActionTransformPipeline( + tokenizer_config=None, + pad_keys=[], + max_action_dim=64, + append_action_caption_semantics=True, + append_camera_role_to_caption=True, + enable_mode_specific_prompt=True, + ) + + result = pipeline(data, "256", action_valid_mask=valid_mask) + + head, left, right = result["ai_caption"] + assert head.startswith("The view is from the camera mounted on the robot head.\n\n") + assert ( + "The action targets are the head ego motion, mapped to the ego-pose." + if head_mapped + else "There is no action target." + ) in head + assert "wrist" not in head + for caption, side, other_side in ((left, "left", "right"), (right, "right", "left")): + assert caption.startswith(f"The view is from the camera mounted on the {side} wrist.\n\n") + assert ( + f"The action targets are the {side} wrist-mounted camera motion and gripper openness, " + f"mapped to the {side} wrist-pose and {side}-openness components." + ) in caption + assert other_side not in caption and "ego-pose" not in caption + for caption in (head, left, right): + assert ("The goal is" in caption) is (mode != "inverse_dynamics") + assert "2.0 seconds long" in caption + assert attributes["subject_postfix"] not in caption + assert bool(result["action_valid_mask"][:9].all()) is head_mapped + if not head_mapped: + assert not bool(result["action_raw"][:16].any()) + assert set(range(16)) <= set(result["sequence_plan"].condition_frame_indexes_action) + assert "action_caption_attributes" not in result + assert attributes["subject_postfix"] == CaptionSubjectMappings["ego_and_dual_wrist_cameras"] @pytest.mark.L0 @@ -317,6 +528,44 @@ def test_action_transform_pipeline_preserves_explicit_action_valid_mask() -> Non torch.testing.assert_close(result["action_processing_record"].action_valid_mask, valid_mask) +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("video_frames", [1, 17]) +@pytest.mark.parametrize("history_steps", [0, 2]) +def test_policy_preserves_action_horizon_history_and_prompt_duration(video_frames: int, history_steps: int) -> None: + pipeline = ActionTransformPipeline( + pad_keys=[], + max_action_dim=4, + append_viewpoint_info=False, + append_resolution_info=False, + format_prompt_as_json=True, + ) + video = torch.arange(3 * video_frames * 4 * 6).reshape(3, video_frames, 4, 6).float() + action = torch.arange(32).reshape(16, 2).float() + history = torch.arange(history_steps * 2).reshape(history_steps, 2).float() + data = { + "ai_caption": "Move the camera.", + "video": video, + "video_num_frames": 17, + "action": action, + "history_action": history, + "conditioning_fps": torch.tensor(8), + "mode": "policy", + } + + result = pipeline(data, resolution="256") + + assert result["video"].shape == (3, 1, 4, 6) + torch.testing.assert_close(result["video"], video[:, :1]) + assert result["ai_caption"]["duration"] == "2s" + assert result["action"].shape == (history_steps + 16, 4) + torch.testing.assert_close(result["action_raw"], torch.cat((history, action))) + assert result["sequence_plan"].condition_frame_indexes_vision == [0] + assert result["sequence_plan"].condition_frame_indexes_action == list(range(history_steps)) + assert result["sequence_plan"].action_start_frame_offset == 1 - history_steps + assert "video_num_frames" not in result + + @pytest.mark.L0 def test_action_transform_pipeline_supports_compact_streaming_video() -> None: pipeline = ActionTransformPipeline( diff --git a/cosmos_framework/data/generator/action/utils/viewpoint_utils.py b/cosmos_framework/data/generator/action/utils/viewpoint_utils.py index a3c12c980..f47d329ea 100644 --- a/cosmos_framework/data/generator/action/utils/viewpoint_utils.py +++ b/cosmos_framework/data/generator/action/utils/viewpoint_utils.py @@ -17,14 +17,24 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log - -Viewpoint = Literal["ego_view", "third_person_view", "wrist_view", "concat_view", "top_down_2d_view", "video_game_view"] +from cosmos_framework.data.generator.action.action_caption_attribute_adapter import ViewCompositionMappings + +Viewpoint = Literal[ + "ego_view", + "third_person_view", + "wrist_view", + "concat_view", + "separate_camera_views", + "top_down_2d_view", + "video_game_view", +] DEFAULT_VIEWPOINT_TEMPLATES: dict[str, str] = { "ego_view": "This video is captured from a first-person perspective looking at the scene.", "third_person_view": "This video is captured from a third-person perspective looking towards the agent from the front.", "wrist_view": "This video is captured from a wrist-mounted camera.", "concat_view": "This video contains concatenated views from multiple camera perspectives.", + "separate_camera_views": ViewCompositionMappings["separate_camera_views"], "top_down_2d_view": "This video is captured from a top-down view of a flat 2D scene.", "video_game_view": "This video is captured from an in-game camera perspective in a virtual environment.", } diff --git a/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py b/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py index b303910fb..0ffd75ca4 100644 --- a/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py +++ b/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py @@ -285,6 +285,8 @@ class InterleavedMediaResizeByMaxPixels(Augmentor): input_keys: List with the single key holding the media dict. Defaults to ``["media_list"]``. max_pixels: Maximum total pixel area per image after resize. padding_divisor: Dimension alignment divisor (both width and height become multiples of it). + interpolation_overrides: Map media keys to Pillow resampling names (e.g. "nearest" or "bilinear"). + Unlisted keys use the default Lanczos interpolation. args: Additional arguments passed to the parent class. """ @@ -294,11 +296,15 @@ def __init__( max_pixels: int = DEFAULT_MAX_PIXELS, padding_divisor: int = 16, args: Optional[dict] = None, + interpolation_overrides: dict[str, str] | None = None, ) -> None: input_keys = input_keys or ["media_list"] super().__init__(input_keys, None, args) self.max_pixels = max_pixels self.padding_divisor = padding_divisor + self.interpolation_overrides: dict[str, Image.Resampling] = { + key: Image.Resampling[mode.upper()] for key, mode in (interpolation_overrides or {}).items() + } def _compute_target_size(self, width: int, height: int) -> tuple[int, int]: """Scale to fit within max_pixels, then align down to padding_divisor.""" @@ -309,17 +315,19 @@ def _compute_target_size(self, width: int, height: int) -> tuple[int, int]: padding_constant=self.padding_divisor, ) - def _resize_image(self, img: Image.Image) -> Image.Image: + def _resize_image(self, img: Image.Image, resample: Image.Resampling = Image.Resampling.LANCZOS) -> Image.Image: w, h = img.size target_w, target_h = self._compute_target_size(w, h) if (target_w, target_h) != (w, h): - img = img.resize((target_w, target_h), Image.Resampling.LANCZOS) + img = img.resize((target_w, target_h), resample) return img - def _process_media(self, media) -> Image.Image | list: + def _process_media( + self, media: Image.Image | list[Image.Image], resample: Image.Resampling = Image.Resampling.LANCZOS + ) -> Image.Image | list[Image.Image]: if isinstance(media, list): - return [self._resize_image(frame) for frame in media] - return self._resize_image(media) + return [self._resize_image(frame, resample) for frame in media] + return self._resize_image(media, resample) def __call__(self, data_dict: Dict) -> Optional[Dict]: assert len(self.input_keys) == 1, ( @@ -333,7 +341,8 @@ def __call__(self, data_dict: Dict) -> Optional[Dict]: diffusion_media_content = {} for key, media in media_list.items(): - diffusion_media_content[key] = self._process_media(media) + resample = self.interpolation_overrides.get(key, Image.Resampling.LANCZOS) + diffusion_media_content[key] = self._process_media(media, resample) data_dict["diffusion_media_list"] = diffusion_media_content return data_dict diff --git a/cosmos_framework/data/generator/augmentors/multi_reference_bbox_condition.py b/cosmos_framework/data/generator/augmentors/multi_reference_bbox_condition.py index c8f687c45..f3d565de8 100644 --- a/cosmos_framework/data/generator/augmentors/multi_reference_bbox_condition.py +++ b/cosmos_framework/data/generator/augmentors/multi_reference_bbox_condition.py @@ -116,7 +116,7 @@ def _draw_reference_number(image: Image.Image, position: tuple[int, int], index: draw.text((x + margin - left, y + margin - top), text, fill=text_color, font=font) -def _format_coordinate_box(box: BoundingBox, coordinate_format: BBoxCoordinateFormat) -> str: +def format_coordinate_box(box: BoundingBox, coordinate_format: BBoxCoordinateFormat) -> str: """Round half up to thousandths while keeping nonempty integer rectangles.""" if coordinate_format == "normalized_01": return "[" + ", ".join(f"{coordinate:.6f}" for coordinate in box) + "]" @@ -173,7 +173,7 @@ def build_bbox_condition( boxes = reference_bboxes.get(key, []) if not boxes: continue - formatted = [_format_coordinate_box(box, coordinate_format) for box in boxes] + formatted = [format_coordinate_box(box, coordinate_format) for box in boxes] coordinates = formatted[0] if len(formatted) == 1 else "[" + ", ".join(formatted) + "]" lines.append(f": {coordinates}") return reference_images, prompt + "\n\n" + "\n".join(lines) diff --git a/cosmos_framework/data/generator/augmentors/multi_reference_subject_layout.py b/cosmos_framework/data/generator/augmentors/multi_reference_subject_layout.py new file mode 100644 index 000000000..bd0de4b74 --- /dev/null +++ b/cosmos_framework/data/generator/augmentors/multi_reference_subject_layout.py @@ -0,0 +1,490 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Subject-bound layout conditioning for the unified multi-reference annotation. + +The unified annotation prompt has three sections:: + + [Entities] + : the woman from , wearing the dress from + : a small brown dog (text-only, no reference) + + [Scene] The overall scene is ... from , ... + [Pose] ... (optional) + [Style] ... (optional) + + [Composition] + sits beside ... + +Boxes are bound to subjects, not to reference images: ``subject_bboxes`` is +``[{"subject": k, "boxes": [[x1, y1, x2, y2], ...]}]`` in target-normalized xyxy. +Subjects without a box (text-only or ungrounded) are absent from the layout. + +The condition adds one ``[Layout]`` line right after ``[Scene]`` (the position +the annotation format renders it at). ``rendered`` mode also produces a white +layout image with one numbered outline per subject box; it becomes reference +```` and is appended after reference shuffling, so its tag is outside +the shuffled range and never renumbered. ``rendered_paste`` mode additionally +pastes each subject's content and worn-item references into its boxes (tiled in +order, aspect ratio kept) under the outlines. Badges show subject numbers, which shuffling never +changes. References themselves are not modified. +""" + +from __future__ import annotations + +import json +import math +import re +from collections import Counter +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Literal + +from PIL import Image, ImageColor, ImageDraw, ImageFont + +from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor +from cosmos_framework.utils import log +from cosmos_framework.data.generator.augmentors.multi_reference_bbox_condition import ( + REFERENCE_BBOX_COLORS, + BBoxCoordinateFormat, + format_coordinate_box, +) + +BoundingBox = tuple[float, float, float, float] +SubjectLayoutMode = Literal["none", "rendered", "rendered_paste", "coordinates"] +SUBJECT_LAYOUT_MODES: tuple[str, ...] = ("none", "rendered", "rendered_paste", "coordinates") +# Modes that append a layout image as the last reference. +IMAGE_LAYOUT_MODES: tuple[str, ...] = ("rendered", "rendered_paste") + +# Same palette as reference-numbered bbox conditioning, indexed by subject number. +SUBJECT_COLORS: tuple[str, ...] = REFERENCE_BBOX_COLORS +SUBJECT_LAYOUT_IMAGE_KEY = "subject_layout_image" + +RENDERED_LAYOUT_NOTE = "Each numbered box marks where the subject with that number appears." +PASTED_LAYOUT_NOTE = ( + "Each numbered box marks where the subject with that number appears and shows the reference images of that subject." +) +_SCENE_PREFIX = "[Scene]" +# Lines rendered after [Scene] in the unified format; the [Layout] line goes before the first of them. +_AFTER_SCENE_MARKERS: tuple[str, ...] = ("\n[Pose] ", "\n[Style] ", "\n\n[Composition]") + + +@dataclass(frozen=True) +class SubjectBoxes: + subject: int + boxes: tuple[BoundingBox, ...] + + +def parse_subject_bboxes(value: object) -> tuple[SubjectBoxes, ...]: + """Validate the ``subject_bboxes`` payload (JSON text or decoded list). + + ``None`` and ``[]`` both mean no layout. Subjects with no boxes are dropped. + """ + if value is None: + return () + if isinstance(value, (str, bytes)): + value = json.loads(value) + if not isinstance(value, list): + raise ValueError("subject_bboxes must be a list.") + seen: set[int] = set() + result: list[SubjectBoxes] = [] + for entry in value: + if not isinstance(entry, dict): + raise ValueError("Each subject bbox entry must be an object.") + subject = entry.get("subject") + if isinstance(subject, bool) or not isinstance(subject, int) or subject < 1: + raise ValueError(f"Subject number must be a positive integer, got {subject!r}.") + if subject in seen: + raise ValueError(f"Duplicate subject {subject} in subject_bboxes.") + seen.add(subject) + raw_boxes = entry.get("boxes") + if not isinstance(raw_boxes, list): + raise ValueError(f"Boxes for subject {subject} must be a list.") + boxes: list[BoundingBox] = [] + for raw_box in raw_boxes: + if not isinstance(raw_box, (list, tuple)) or len(raw_box) != 4: + raise ValueError(f"Each box for subject {subject} must contain four xyxy coordinates.") + if any(isinstance(c, bool) or not isinstance(c, (int, float)) or not math.isfinite(c) for c in raw_box): + raise ValueError(f"Coordinates for subject {subject} must be finite numbers.") + x1, y1, x2, y2 = (float(c) for c in raw_box) + if not (0 <= x1 < x2 <= 1 and 0 <= y1 < y2 <= 1): + raise ValueError(f"Boxes for subject {subject} must be nonempty xyxy rectangles in [0, 1].") + boxes.append((x1, y1, x2, y2)) + if boxes: + result.append(SubjectBoxes(subject=subject, boxes=tuple(boxes))) + return tuple(result) + + +def drop_shared_boxes(subjects: tuple[SubjectBoxes, ...]) -> tuple[SubjectBoxes, ...]: + """Remove boxes assigned to more than one subject; subjects left without boxes are dropped. + + Boxes are copied from reference geometry, so two subjects get the same box when they + come from one reference image (several things in one crop, or repeated copies of one + object) or when upstream gave two references the same box. Such a box does not say + where each subject is. + """ + counts = Counter(box for item in subjects for box in set(item.boxes)) + result: list[SubjectBoxes] = [] + for item in subjects: + boxes = tuple(box for box in item.boxes if counts[box] == 1) + if boxes: + result.append(SubjectBoxes(subject=item.subject, boxes=boxes)) + return tuple(result) + + +_SUBJECT_TAG_RE = re.compile(r"") + + +def defined_subjects(prompt: str) -> set[int]: + """Subject numbers defined in the ``[Entities]`` section (``:`` line starts).""" + entities = prompt.split("\n\n", 1)[0] + return { + int(match.group(1)) + for line in entities.splitlines() + if (match := _SUBJECT_TAG_RE.match(line)) is not None and line[match.end() :].startswith(":") + } + + +def insert_layout_line(prompt: str, line: str) -> str: + """Insert a ``[Layout]`` line directly after the ``[Scene]`` line of a unified annotation.""" + scene_start = prompt.find(_SCENE_PREFIX) + if scene_start < 0: + raise ValueError("Prompt has no [Scene] line; it is not a unified annotation.") + if "\n[Layout] " in prompt: + raise ValueError("Prompt already has a [Layout] line.") + positions = [position for marker in _AFTER_SCENE_MARKERS if (position := prompt.find(marker, scene_start)) >= 0] + if not positions: + raise ValueError("Prompt has no [Composition] section after [Scene].") + insert_at = min(positions) + return prompt[:insert_at] + "\n" + line + prompt[insert_at:] + + +def rendered_layout_line(layout_image_index: int, note: str = RENDERED_LAYOUT_NOTE) -> str: + return ( + f"[Layout] The overall layout of all the mentioned entities is provided in . {note}" + ) + + +def coordinates_layout_line( + subjects: tuple[SubjectBoxes, ...], coordinate_format: BBoxCoordinateFormat = "integer_1000" +) -> str: + """Boxes as text, formatted like reference-numbered bbox coordinates (integers in [0, 1000] by default).""" + parts = [] + for item in subjects: + boxes = ", ".join(format_coordinate_box(box, coordinate_format) for box in item.boxes) + parts.append(f" {boxes}") + scale = "integers in [0, 1000]" if coordinate_format == "integer_1000" else "values in [0, 1]" + return ( + f"[Layout] The positions of the subjects are given as [x1, y1, x2, y2] boxes normalized to {scale}: " + + "; ".join(parts) + + "." + ) + + +Rect = tuple[int, int, int, int] + + +def _overlap(a: Rect, b: Rect) -> int: + return max(0, min(a[2], b[2]) - max(a[0], b[0])) * max(0, min(a[3], b[3]) - max(a[1], b[1])) + + +def _label_candidates(box: Rect, width: int, height: int) -> list[tuple[int, int]]: + """Badge anchors around a box, in preference order: corners inside, then just outside.""" + x1, y1, x2, y2 = box + return [ + (x1, y1), + (x2 - width, y1), + (x1, y1 - height), + (x2 - width, y1 - height), + (x1, y2 - height), + (x2 - width, y2 - height), + (x1, y2), + ((x1 + x2 - width) // 2, (y1 + y2 - height) // 2), + ] + + +def place_labels(boxes: list[Rect], size: tuple[int, int], label_size: tuple[int, int]) -> list[Rect]: + """Greedy deterministic badge placement: the first on-canvas anchor that overlaps no earlier badge, + else the anchor with the least overlap.""" + canvas_w, canvas_h = size + width, height = label_size + placed: list[Rect] = [] + for box in boxes: + best: Rect | None = None + best_cost: int | None = None + for anchor_x, anchor_y in _label_candidates(box, width, height): + x = max(0, min(anchor_x, canvas_w - width)) + y = max(0, min(anchor_y, canvas_h - height)) + rect = (x, y, x + width, y + height) + cost = sum(_overlap(rect, other) for other in placed) + if cost == 0: + best = rect + break + if best_cost is None or cost < best_cost: + best, best_cost = rect, cost + assert best is not None + placed.append(best) + return placed + + +def _fit(image: Image.Image, width: int, height: int) -> Image.Image: + scale = min(width / image.width, height / image.height) + return image.convert("RGB").resize( + (max(1, round(image.width * scale)), max(1, round(image.height * scale))), Image.Resampling.BICUBIC + ) + + +def _grid_columns(images: Sequence[Image.Image], width: int, height: int) -> int: + """Column count whose grid shows the largest total pasted area (ties prefer fewer columns).""" + best_columns, best_area = 1, -1.0 + for columns in range(1, len(images) + 1): + rows = math.ceil(len(images) / columns) + cell_width, cell_height = width / columns, height / rows + area = sum( + image.width * image.height * min(cell_width / image.width, cell_height / image.height) ** 2 + for image in images + ) + if area > best_area: + best_columns, best_area = columns, area + return best_columns + + +# A subject's references avoid boxes pasted over them unless the free part is smaller than this share of the box. +MIN_FREE_PASTE_AREA = 0.25 + + +def _largest_free_rect(box: Rect, obstacles: list[Rect]) -> Rect: + """Largest axis-aligned rectangle inside ``box`` that overlaps none of ``obstacles``.""" + x1, y1, x2, y2 = box + xs = sorted({x1, x2, *(min(max(x, x1), x2) for o in obstacles for x in (o[0], o[2]))}) + best, best_area = (x1, y1, x1, y1), 0 + for i, left in enumerate(xs): + for right in xs[i + 1 :]: + spans = sorted( + (max(o[1], y1), min(o[3], y2)) + for o in obstacles + if o[0] < right and o[2] > left and o[1] < y2 and o[3] > y1 + ) + cursor = y1 + for top, bottom in [*spans, (y2, y2)]: + if (right - left) * (top - cursor) > best_area: + best, best_area = (left, cursor, right, top), (right - left) * (top - cursor) + cursor = max(cursor, bottom) + return best + + +def _tile(canvas: Image.Image, images: Sequence[Image.Image], area: Rect) -> None: + """Tile ``images`` in order, row by row, inside ``area``; each is fitted and centered in its cell.""" + left, top = area[0], area[1] + area_width, area_height = max(1, area[2] - area[0]), max(1, area[3] - area[1]) + columns = _grid_columns(images, area_width, area_height) + rows = math.ceil(len(images) / columns) + for index, image in enumerate(images): + row, column = divmod(index, columns) + cell_left = left + column * area_width // columns + cell_top = top + row * area_height // rows + cell_width = left + (column + 1) * area_width // columns - cell_left + cell_height = top + (row + 1) * area_height // rows - cell_top + fitted = _fit(image, max(1, cell_width), max(1, cell_height)) + canvas.paste( + fitted, (cell_left + (cell_width - fitted.width) // 2, cell_top + (cell_height - fitted.height) // 2) + ) + + +def paste_subject_references( + canvas: Image.Image, subjects: tuple[SubjectBoxes, ...], subject_images: Mapping[int, Sequence[Image.Image]] +) -> None: + """Paste each subject's references into its boxes. + + A subject's references (content images first, then worn-item references) are tiled + in order, row by row, in the grid that shows them largest; each keeps its aspect ratio + and is centered in its cell. Larger boxes are pasted first, so smaller ones (e.g. held + objects) stay visible on top. The tiles use the largest part of the box that no later + (smaller) box of another subject covers, so no reference is hidden; when that part is + under ``MIN_FREE_PASTE_AREA`` of the box, the whole box is used. Subjects without images + keep an empty box. + """ + width, height = canvas.size + items: list[tuple[int, Rect]] = [ + (item.subject, (round(x1 * width), round(y1 * height), round(x2 * width), round(y2 * height))) + for item in subjects + if subject_images.get(item.subject) + for x1, y1, x2, y2 in item.boxes + ] + items.sort(key=lambda item: -(item[1][2] - item[1][0]) * (item[1][3] - item[1][1])) + for position, (subject, box) in enumerate(items): + later = [other for other_subject, other in items[position + 1 :] if other_subject != subject] + free = _largest_free_rect(box, later) + box_area = max(1, (box[2] - box[0]) * (box[3] - box[1])) + use_free = (free[2] - free[0]) * (free[3] - free[1]) >= MIN_FREE_PASTE_AREA * box_area + _tile(canvas, subject_images[subject], free if use_free else box) + + +def render_subject_layout( + subjects: tuple[SubjectBoxes, ...], + size: tuple[int, int], + subject_images: Mapping[int, Sequence[Image.Image]] | None = None, +) -> Image.Image: + """White canvas at the target size with one numbered outline per subject box. + + With ``subject_images`` (subject -> its references), the references are pasted into + their boxes first and the outlines and badges are drawn on top. + """ + width, height = size + if width <= 0 or height <= 0: + raise ValueError("Layout dimensions must be positive.") + canvas = Image.new("RGB", size, "white") + if subject_images: + paste_subject_references(canvas, subjects, subject_images) + draw = ImageDraw.Draw(canvas) + line_width = max(2, round(min(size) / 256)) + font_size = max(12, round(min(size) / 28)) + font = ImageFont.load_default(size=font_size) + margin = max(2, font_size // 5) + + items: list[tuple[int, Rect]] = [] + for item in subjects: + for x1, y1, x2, y2 in item.boxes: + rect = ( + min(width - 1, round(x1 * width)), + min(height - 1, round(y1 * height)), + min(width - 1, round(x2 * width)), + min(height - 1, round(y2 * height)), + ) + items.append((item.subject, rect)) + # Every badge has the size of the widest label. + widest = max((str(subject) for subject, _ in items), key=len, default="0") + left, top, right, bottom = draw.textbbox((0, 0), widest, font=font) + label_size = (math.ceil(right - left) + 2 * margin, math.ceil(bottom - top) + 2 * margin) + + for subject, rect in items: + draw.rectangle(rect, outline=SUBJECT_COLORS[(subject - 1) % len(SUBJECT_COLORS)], width=line_width) + for (subject, _), badge in zip(items, place_labels([rect for _, rect in items], size, label_size), strict=True): + color = SUBJECT_COLORS[(subject - 1) % len(SUBJECT_COLORS)] + draw.rectangle((badge[0], badge[1], badge[2] - 1, badge[3] - 1), fill=color) + red, green, blue = ImageColor.getrgb(color)[:3] + text_color = "white" if 299 * red + 587 * green + 114 * blue < 140000 else "black" + text = str(subject) + t_left, t_top, t_right, _ = draw.textbbox((0, 0), text, font=font) + x = badge[0] + (label_size[0] - (t_right - t_left)) // 2 - t_left + draw.text((x, badge[1] + margin - t_top), text, fill=text_color, font=font) + return canvas + + +def build_subject_layout_condition( + prompt: str, + subject_bboxes: object, + num_references: int, + target_size: tuple[int, int], + mode: SubjectLayoutMode, + *, + coordinate_format: BBoxCoordinateFormat = "integer_1000", + subject_images: Mapping[int, Sequence[Image.Image]] | None = None, +) -> tuple[str, Image.Image | None]: + """Return the conditioned prompt and, in image modes, the layout image for ````. + + ``rendered_paste`` pastes ``subject_images`` (subject -> its content references, then + its worn-item references) into the boxes; subjects missing from it keep an empty box. + + Without boxes (NULL or empty geometry, or only boxes shared by several subjects) + the prompt is returned unchanged, so a layout recipe can mix in populations that + have no geometry as plain samples. + """ + if mode not in SUBJECT_LAYOUT_MODES: + raise ValueError(f"Unknown subject layout mode: {mode!r}.") + if coordinate_format not in ("integer_1000", "normalized_01"): + raise ValueError(f"Unknown bbox coordinate format: {coordinate_format!r}.") + if mode == "none": + return prompt, None + subjects = parse_subject_bboxes(subject_bboxes) + undefined = {item.subject for item in subjects} - defined_subjects(prompt) + if undefined: + raise ValueError(f"Boxes refer to subjects {sorted(undefined)} that the prompt does not define.") + subjects = drop_shared_boxes(subjects) + if not subjects: + return prompt, None + if mode == "coordinates": + return insert_layout_line(prompt, coordinates_layout_line(subjects, coordinate_format)), None + if mode == "rendered_paste": + layout_image = render_subject_layout(subjects, target_size, subject_images or {}) + return insert_layout_line(prompt, rendered_layout_line(num_references + 1, PASTED_LAYOUT_NOTE)), layout_image + layout_image = render_subject_layout(subjects, target_size) + return insert_layout_line(prompt, rendered_layout_line(num_references + 1)), layout_image + + +def subject_reference_images( + subject_paste_images: Mapping[int, Sequence[int]], references: list[Image.Image] +) -> dict[int, list[Image.Image]]: + """Subject -> its reference images, from 1-based reference indices.""" + images: dict[int, list[Image.Image]] = {} + for subject, indices in subject_paste_images.items(): + for index in indices: + if not 1 <= index <= len(references): + raise ValueError( + f"Subject {subject} refers to but there are {len(references)} references." + ) + images[subject] = [references[index - 1] for index in indices] + return images + + +class BuildSubjectLayoutCondition(Augmentor): + """Add the subject layout to the prompt before references are shuffled. + + Reads ``annotations["subject_bboxes"]``, ``source_image`` and ``target_image`` + (after resize); ``rendered_paste`` also reads ``annotations["subject_paste_images"]`` + (subject -> 1-based indices of its content and worn-item references, in the unshuffled order). + The layout image is kept under ``subject_layout_image`` and appended by + :class:`AppendSubjectLayoutImage` after the shuffle. + """ + + def __init__( + self, + mode: SubjectLayoutMode, + coordinate_format: BBoxCoordinateFormat = "integer_1000", + input_keys: list[str] | None = None, + ) -> None: + super().__init__(input_keys or []) + if mode not in ("rendered", "rendered_paste", "coordinates"): + raise ValueError("The subject layout augmentor requires rendered, rendered_paste or coordinates mode.") + if coordinate_format not in ("integer_1000", "normalized_01"): + raise ValueError(f"Unknown bbox coordinate format: {coordinate_format!r}.") + self.mode: SubjectLayoutMode = mode + self.coordinate_format: BBoxCoordinateFormat = coordinate_format + + def __call__(self, data_dict: dict) -> dict | None: + try: + references = data_dict["source_image"] + subject_images = None + if self.mode == "rendered_paste": + subject_images = subject_reference_images(data_dict["annotations"]["subject_paste_images"], references) + prompt, layout_image = build_subject_layout_condition( + prompt=data_dict["editing_instruction"], + subject_bboxes=data_dict["annotations"].get("subject_bboxes"), + num_references=len(references), + target_size=data_dict["target_image"].size, + mode=self.mode, + coordinate_format=self.coordinate_format, + subject_images=subject_images, + ) + except (KeyError, ValueError) as error: + log.warning( + f"Skipping subject layout sample {data_dict.get('__key__', 'unknown')!r}: {error}", rank0_only=False + ) + return None + data_dict["editing_instruction"] = prompt + if layout_image is not None: + data_dict[SUBJECT_LAYOUT_IMAGE_KEY] = layout_image + return data_dict + + +class AppendSubjectLayoutImage(Augmentor): + """Append the rendered layout as the last reference, after reference shuffling.""" + + def __init__(self, input_keys: list[str] | None = None) -> None: + super().__init__(input_keys or []) + + def __call__(self, data_dict: dict) -> dict | None: + layout_image = data_dict.pop(SUBJECT_LAYOUT_IMAGE_KEY, None) + if layout_image is not None: + data_dict["source_image"] = [*data_dict["source_image"], layout_image] + return data_dict diff --git a/cosmos_framework/data/generator/augmentors/multi_reference_subject_layout_test.py b/cosmos_framework/data/generator/augmentors/multi_reference_subject_layout_test.py new file mode 100644 index 000000000..c8ac8787e --- /dev/null +++ b/cosmos_framework/data/generator/augmentors/multi_reference_subject_layout_test.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Tests for subject-bound layout conditioning of unified multi-reference prompts.""" + +import json +import random + +import pytest +from PIL import Image + +from cosmos_framework.data.generator.augmentors.multi_reference_subject_layout import ( + SUBJECT_COLORS, + AppendSubjectLayoutImage, + BuildSubjectLayoutCondition, + build_subject_layout_condition, + defined_subjects, + insert_layout_line, + parse_subject_bboxes, + place_labels, +) +from cosmos_framework.data.generator.augmentors.multi_reference_transform import ReorderReferenceImages + +PROMPT = ( + "[Entities]\n" + ": the woman from , wearing the red dress from .\n" + ": a small brown dog.\n\n" + "[Scene] The overall scene is a reading nook from , with warm light.\n" + "[Style] The overall style of the image is provided in , a watercolor.\n\n" + "[Composition]\n" + " sits on the rug beside ." +) +BOXES = json.dumps([{"subject": 1, "boxes": [[0.1, 0.2, 0.5, 0.9]]}, {"subject": 2, "boxes": [[0.6, 0.6, 0.9, 0.95]]}]) + + +def _sample(prompt: str = PROMPT, boxes: str | None = BOXES, num_references: int = 4) -> dict: + return { + "__key__": "view-0", + "annotations": {"subject_bboxes": boxes, "subject_paste_images": {1: [1]}}, + "source_image": [Image.new("RGB", (32, 32), (index * 60, 0, 0)) for index in range(num_references)], + "target_image": Image.new("RGB", (200, 100)), + "editing_instruction": prompt, + } + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_layout_line_goes_right_after_scene() -> None: + text = insert_layout_line(PROMPT, "[Layout] X.") + assert "[Scene] The overall scene is a reading nook from , with warm light.\n[Layout] X.\n[Style]" in text + without_style = PROMPT.replace( + "\n[Style] The overall style of the image is provided in , a watercolor.", "" + ) + assert "warm light.\n[Layout] X.\n\n[Composition]" in insert_layout_line(without_style, "[Layout] X.") + with pytest.raises(ValueError, match="no \\[Scene\\]"): + insert_layout_line(" a cat", "[Layout] X.") + with pytest.raises(ValueError, match="already"): + insert_layout_line(text, "[Layout] Y.") + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_rendered_mode_adds_layout_image_as_next_reference() -> None: + prompt, layout = build_subject_layout_condition(PROMPT, BOXES, 4, (200, 100), "rendered") + assert "\n[Layout] The overall layout of all the mentioned entities is provided in ." in prompt + assert layout is not None and layout.size == (200, 100) + # Subject 1's box outline uses the first palette color. + assert layout.getpixel((20, 60)) == Image.new("RGB", (1, 1), SUBJECT_COLORS[0]).getpixel((0, 0)) + assert layout.getpixel((70, 60)) == (255, 255, 255) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_rendered_paste_mode_pastes_subject_reference_under_outlines() -> None: + reference = Image.new("RGB", (40, 80), (0, 0, 255)) + prompt, layout = build_subject_layout_condition( + PROMPT, BOXES, 4, (200, 100), "rendered_paste", subject_images={1: [reference]} + ) + assert ( + "provided in . Each numbered box marks where the subject with that number appears and shows " + "the reference images of that subject." in prompt + ) + assert layout is not None + # Subject 1's box is 80x70 px; the 40x80 reference is fitted to 35x70 and centered in it. + assert layout.getpixel((60, 55)) == (0, 0, 255) + assert layout.getpixel((30, 55)) == (255, 255, 255) + assert layout.getpixel((20, 60)) == Image.new("RGB", (1, 1), SUBJECT_COLORS[0]).getpixel((0, 0)) + # Subject 2 has no content image: an empty outlined box. + assert layout.getpixel((150, 80)) == (255, 255, 255) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_rendered_paste_tiles_a_subjects_references_in_order() -> None: + person, dress = Image.new("RGB", (40, 80), (0, 0, 255)), Image.new("RGB", (40, 80), (0, 255, 0)) + boxes = json.dumps([{"subject": 1, "boxes": [[0.1, 0.2, 0.9, 0.9]]}]) + _, layout = build_subject_layout_condition( + PROMPT, boxes, 4, (200, 100), "rendered_paste", subject_images={1: [person, dress]} + ) + assert layout is not None + # A wide 160x70 box becomes two 80x70 cells; the tall references sit side by side, content first. + assert layout.getpixel((60, 55)) == (0, 0, 255) + assert layout.getpixel((140, 55)) == (0, 255, 0) + assert layout.getpixel((100, 55)) == (255, 255, 255) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_rendered_paste_keeps_references_clear_of_smaller_boxes_pasted_on_top() -> None: + person, bag = Image.new("RGB", (40, 40), (0, 0, 255)), Image.new("RGB", (40, 40), (0, 255, 0)) + # The bag box covers the left half of the person box, so the person goes into the right half. + boxes = json.dumps( + [{"subject": 1, "boxes": [[0.0, 0.0, 1.0, 1.0]]}, {"subject": 2, "boxes": [[0.0, 0.0, 0.5, 1.0]]}] + ) + _, layout = build_subject_layout_condition( + PROMPT, boxes, 4, (200, 100), "rendered_paste", subject_images={1: [person], 2: [bag]} + ) + assert layout is not None + assert layout.getpixel((150, 50)) == (0, 0, 255) + assert layout.getpixel((50, 50)) == (0, 255, 0) + # When less than a quarter of the box is free, the whole box is used. + covered = json.dumps( + [{"subject": 1, "boxes": [[0.0, 0.0, 1.0, 1.0]]}, {"subject": 2, "boxes": [[0.0, 0.0, 0.9, 1.0]]}] + ) + _, layout = build_subject_layout_condition( + PROMPT, covered, 4, (200, 100), "rendered_paste", subject_images={1: [person], 2: [bag]} + ) + assert layout is not None + # The person fills the whole box (x 50..150); the bag (x 40..140) covers all but its right edge. + assert layout.getpixel((147, 50)) == (0, 0, 255) + assert layout.getpixel((190, 50)) == (255, 255, 255) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_coordinates_mode_lists_integer_boxes_per_subject_without_image() -> None: + prompt, layout = build_subject_layout_condition(PROMPT, BOXES, 4, (200, 100), "coordinates") + assert layout is None + assert ( + "[Layout] The positions of the subjects are given as [x1, y1, x2, y2] boxes normalized to integers in " + "[0, 1000]: [100, 200, 500, 900]; [600, 600, 900, 950]." in prompt + ) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_coordinates_mode_keeps_legacy_normalized_format() -> None: + prompt, _ = build_subject_layout_condition( + PROMPT, BOXES, 4, (200, 100), "coordinates", coordinate_format="normalized_01" + ) + assert "boxes normalized to values in [0, 1]: [0.100000, 0.200000, 0.500000, 0.900000];" in prompt + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("boxes", [None, "[]", json.dumps([{"subject": 1, "boxes": []}])]) +def test_missing_geometry_keeps_plain_prompt(boxes: str | None) -> None: + assert build_subject_layout_condition(PROMPT, boxes, 4, (200, 100), "rendered") == (PROMPT, None) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_boxes_shared_by_several_subjects_are_dropped() -> None: + prompt = PROMPT.replace( + ": a small brown dog.", ": a small brown dog.\n: a red ball." + ) + shared, own = [0.1, 0.2, 0.5, 0.9], [0.6, 0.6, 0.9, 0.95] + boxes = json.dumps( + [{"subject": 1, "boxes": [shared]}, {"subject": 2, "boxes": [shared, own]}, {"subject": 3, "boxes": [shared]}] + ) + text, _ = build_subject_layout_condition(prompt, boxes, 4, (200, 100), "coordinates") + assert "[0, 1000]: [600, 600, 900, 950]." in text + only_shared = json.dumps([{"subject": 1, "boxes": [shared]}, {"subject": 2, "boxes": [shared]}]) + assert build_subject_layout_condition(PROMPT, only_shared, 4, (200, 100), "rendered") == (PROMPT, None) + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize( + "boxes", + [ + json.dumps([{"subject": 3, "boxes": [[0.1, 0.1, 0.2, 0.2]]}]), + json.dumps([{"subject": 1, "boxes": [[0.5, 0.1, 0.2, 0.2]]}]), + json.dumps([{"subject": 1, "boxes": [[0.1, 0.1, 0.2]]}]), + json.dumps([{"subject": 1, "boxes": []}, {"subject": 1, "boxes": []}]), + ], +) +def test_invalid_geometry_skips_sample(boxes: str) -> None: + assert BuildSubjectLayoutCondition("rendered")(_sample(boxes=boxes)) is None + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_defined_subjects_reads_entity_lines_only() -> None: + assert defined_subjects(PROMPT) == {1, 2} + assert parse_subject_bboxes(None) == () + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("mode", ["rendered", "rendered_paste"]) +def test_shuffle_renumbers_references_but_not_the_appended_layout(mode: str) -> None: + random.seed(3) + data = BuildSubjectLayoutCondition(mode)(_sample()) + assert data is not None + originals = list(data["source_image"]) + data = ReorderReferenceImages(shuffle_prob=1.0)(data) + data = AppendSubjectLayoutImage()(data) + sources = data["source_image"] + assert len(sources) == 5 and sources[-1].size == (200, 100) + new_position = {id(image): index + 1 for index, image in enumerate(sources[:4])} + prompt = data["editing_instruction"] + assert f"the woman from " in prompt + assert f"red dress from " in prompt + assert f"reading nook from " in prompt + assert "is provided in . Each numbered box" in prompt + assert "subject_layout_image" not in data + # The pasted reference is subject 1's own image, whatever its position after the shuffle. + assert sources[-1].getpixel((60, 55)) == ((0, 0, 0) if mode == "rendered_paste" else (255, 255, 255)) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_rendered_paste_skips_samples_with_out_of_range_content_images() -> None: + sample = _sample() + sample["annotations"]["subject_paste_images"] = {1: [1, 9]} + assert BuildSubjectLayoutCondition("rendered_paste")(sample) is None + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_legacy_img_markers_still_shuffle() -> None: + random.seed(0) + data = {"source_image": ["a", "b", "c"], "editing_instruction": " and in "} + out = ReorderReferenceImages(shuffle_prob=1.0)(data) + order = {name: index + 1 for index, name in enumerate(out["source_image"])} + assert out["editing_instruction"] == f" and in " + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_labels_never_overlap_when_space_allows() -> None: + boxes = [(10, 10, 60, 60), (12, 12, 62, 62), (14, 14, 64, 64)] + badges = place_labels(boxes, (200, 200), (12, 12)) + for index, first in enumerate(badges): + for second in badges[index + 1 :]: + assert min(first[2], second[2]) <= max(first[0], second[0]) or min(first[3], second[3]) <= max( + first[1], second[1] + ) diff --git a/cosmos_framework/data/generator/augmentors/multi_reference_transform.py b/cosmos_framework/data/generator/augmentors/multi_reference_transform.py index 24dbbf9d5..d247187df 100644 --- a/cosmos_framework/data/generator/augmentors/multi_reference_transform.py +++ b/cosmos_framework/data/generator/augmentors/multi_reference_transform.py @@ -216,6 +216,9 @@ class ExtractMultiReferenceConversation(Augmentor): Returns ``None`` when required fields are missing. + ``truncate_excess_references=False`` skips samples above the limit instead of + keeping the first references (prompts that tag every image need all of them). + When ``read_reference_bboxes`` is enabled, also preserve logical input keys and validated target-normalized boxes. Scene and abstract-control references without boxes, including samples with an empty bbox list, remain valid. @@ -231,6 +234,7 @@ def __init__( prompt_variant_keys: tuple[str, ...] | list[str] = _PROMPT_VARIANT_KEYS, args: Optional[dict] = None, read_reference_bboxes: bool = False, + truncate_excess_references: bool = True, ) -> None: super().__init__(input_keys or [], args=args) if max_reference_images <= 0: @@ -241,6 +245,7 @@ def __init__( self.instruction_key = instruction_key self.prompt_variant_keys = tuple(prompt_variant_keys) self.read_reference_bboxes: bool = read_reference_bboxes + self.truncate_excess_references: bool = truncate_excess_references def _resolve_instruction(self, annotation: dict) -> str | None: """Resolve the instruction string, sampling a variant when given a dict. @@ -321,6 +326,14 @@ def __call__(self, data_dict: dict) -> dict | None: return None if len(ordered_input_keys) > self.max_reference_images: + if not self.truncate_excess_references: + log.warning( + f"Skipping sample {data_dict.get('__key__', 'unknown')!r}: " + f"{len(ordered_input_keys)} references exceed the limit {self.max_reference_images}; " + "truncating would invalidate the prompt's image references.", + rank0_only=False, + ) + return None if self.read_reference_bboxes: log.warning( f"Skipping bbox sample {data_dict.get('__key__', 'unknown')!r}: " @@ -431,11 +444,12 @@ def __call__(self, data_dict: dict) -> dict | None: return data_dict -_MARKER_RE: re.Pattern[str] = re.compile(r"") +# Legacy prompts tag references as ````; the unified annotation format uses ````. +_MARKER_RE: re.Pattern[str] = re.compile(r"<(img-|image )(\d+)>") class ReorderReferenceImages(Augmentor): - """Shuffle the reference list and rewrite ```` markers in the instruction. + """Shuffle the reference list and rewrite ```` / ```` markers in the instruction. Operates on the ``source_image`` list and ``editing_instruction`` string produced by :class:`ExtractMultiReferenceConversation`. With probability ``shuffle_prob`` (one @@ -507,9 +521,9 @@ def __call__(self, data_dict: dict) -> dict | None: inv[old_idx] = new_idx def _remap(match: re.Match) -> str: - old_idx = int(match.group(1)) - 1 + old_idx = int(match.group(2)) - 1 if 0 <= old_idx < n: - return f"" + return f"<{match.group(1)}{inv[old_idx] + 1}>" return match.group(0) data_dict[self.source_key] = [source_image[perm[i]] for i in range(n)] diff --git a/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py b/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py index 842bbddc8..185f753aa 100644 --- a/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py @@ -17,19 +17,23 @@ from PIL import Image, UnidentifiedImageError from torchcodec.decoders import AudioDecoder, VideoDecoder -from cosmos_framework.data.generator.processors.qwen3vl_processor import Qwen3VLProcessor -from cosmos_framework.data.generator.reasoner.video_decoder_qwen import VideoTemporalMode, _video_decoder_qwen_func from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log +from cosmos_framework.data.generator.reasoner.video_decoder_qwen import ( + VideoDecodeTimeoutError, + VideoTemporalMode, + _video_decoder_qwen_func, +) +from cosmos_framework.data.generator.processors.qwen3vl_processor import Qwen3VLProcessor from cosmos_framework.utils.generator.source_video_timing import ( SOURCE_VIDEO_TIMING_KEY, require_source_pts_processor, validate_source_video_timing, validate_video_timestamp_mode, ) -from cosmos_framework.utils.generator.torchcodec_video import probe_video from cosmos_framework.utils.generator.video_preprocess import tensor_to_pil_images from cosmos_framework.utils.generator.video_source_metadata import VIDEO_METADATA_KEY +from cosmos_framework.utils.generator.torchcodec_video import probe_video class BytesToMedia(Augmentor): @@ -51,10 +55,11 @@ def __init__( input_key: str = "media", output_key: str = "media", min_fps_thres: int = 4, - max_fps_thres: int = 60, + max_fps_thres: float = 60.0, target_fps: float = 4.0, min_video_token_length: int = 16, max_video_token_length: int = 8192, + max_video_frames: int | None = None, num_threads: int = 0, random_augmentation: bool = False, is_input_pickle_byptes: bool = True, @@ -65,6 +70,8 @@ def __init__( audio_sample_rate: int = 16_000, video_timestamp_mode: str = "qwen_index", video_temporal_mode: VideoTemporalMode = "native", + decoding_timeout: int = 60, + raise_on_decode_timeout: bool = False, ) -> None: """ Args: @@ -82,8 +89,12 @@ def __init__( frame_count_random_range (list[int], optional): Random frame count range. Defaults to None. extract_audio (bool): Whether to decode the audio stream from video containers. audio_sample_rate (int): Target sample rate for decoded mono audio. + decoding_timeout (int): Maximum wait for decoding one chunk of video frames, in seconds. + raise_on_decode_timeout (bool): Keep resource failures out of immutable packing catalogs. """ validate_video_timestamp_mode(video_timestamp_mode) + if decoding_timeout <= 0: + raise ValueError("decoding_timeout must be positive") self.video_timestamp_mode: str = video_timestamp_mode if video_timestamp_mode == "source_pts": require_source_pts_processor(processor) @@ -98,11 +109,14 @@ def __init__( "max_fps_thres": max_fps_thres, "target_fps": target_fps, "min_video_token_length": min_video_token_length, - "max_video_token_length": max_video_token_length, + "max_video_token_length": max_video_token_length, # per sample; _get_decoder_params splits it across videos + "max_video_frames": max_video_frames, "num_threads": num_threads, "random_augmentation": random_augmentation, "frame_count_random_range": frame_count_random_range, "video_temporal_mode": video_temporal_mode, + "decoding_timeout": decoding_timeout, + "raise_on_decode_timeout": raise_on_decode_timeout, } self.is_input_pickle_byptes = is_input_pickle_byptes self.use_start_frame_end_frame = use_start_frame_end_frame @@ -129,7 +143,7 @@ def _is_video_key(self, name: str) -> bool: def _is_audio_key(name: str) -> bool: """Return whether a media key denotes an audio clip.""" name_lower = name.lower() - return "audio" in name_lower or "sound" in name_lower or name_lower.endswith((".wav", ".flac")) + return "audio" in name_lower or "sound" in name_lower or name_lower.endswith((".wav", ".flac", ".mp3")) def _bytes_to_audio_waveform( self, @@ -262,6 +276,10 @@ def _bytes_to_video_frames( result["videos"], channels_first=True ) # 3,T,H,W -> list of PIL images return result + except VideoDecodeTimeoutError: + # A timeout depends on worker load, not on the source row. Packed + # preprocessing must stop rather than seal a different sample set. + raise except Exception as e: if self.video_timestamp_mode == "source_pts": raise ValueError(f"source_pts failed to decode and align video {identifier!r}") from e diff --git a/cosmos_framework/data/generator/augmentors/reasoner/filter_seq_length.py b/cosmos_framework/data/generator/augmentors/reasoner/filter_seq_length.py index 84c36305a..f9cffd0f5 100644 --- a/cosmos_framework/data/generator/augmentors/reasoner/filter_seq_length.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/filter_seq_length.py @@ -3,7 +3,8 @@ """Augmentations to remove keys from the output data_dict""" -from typing import Dict, List, Optional +from collections.abc import Mapping +from typing import Any, List, Optional from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log @@ -32,7 +33,20 @@ def __init__( audio_pad_token: str = AUDIO_PAD_TOKEN, audio_end_token: str = AUDIO_END_TOKEN, drop_over_max_length: bool = False, + complete_sample_max_length_by_source: Mapping[str, int] | None = None, ) -> None: + """Optionally retain complete supervised examples for exact source URL roots. + + Sources in the map use their own limit, drop overlength examples instead + of truncating, and require at least one next-token supervision target. + All other sources retain the existing length and truncation policy. + """ + self.complete_sample_max_length_by_source: dict[str, int] = dict(complete_sample_max_length_by_source or {}) + if any( + not isinstance(root, str) or not root or type(limit) is not int or limit < 2 + for root, limit in self.complete_sample_max_length_by_source.items() + ): + raise ValueError("Complete-sample length policies require non-empty source roots and integer limits >= 2.") self.max_token_length = max_token_length self.processor = processor self.drop_over_max_length: bool = drop_over_max_length @@ -47,8 +61,18 @@ def __init__( audio_end_token=audio_end_token, ).token_ids[1] - def __call__(self, data_dict: Dict) -> Dict: + def __call__(self, data_dict: dict[str, Any]) -> dict[str, Any] | None: input_ids = data_dict["input_ids"] + if self.complete_sample_max_length_by_source: + limit = self.complete_sample_max_length_by_source.get(str(data_dict["__url__"].root)) + if limit is not None: + # Prefix truncation can remove the answer from a long SFT example. + if input_ids.shape[-1] > limit: + return None + labels = data_dict["labels"] + if not bool((labels[1:] != int(data_dict.get("ignore_index", -100))).any()): + return None + return data_dict if input_ids.shape[-1] > self.max_token_length and not self.drop_over_max_length: # check if there is pixel values or pixel value videos in the remaining tokens, if not truncate the input ids input_ids_extra = input_ids[self.max_token_length :] diff --git a/cosmos_framework/data/generator/augmentors/reasoner/format_describe_anything.py b/cosmos_framework/data/generator/augmentors/reasoner/format_describe_anything.py index a955f3e91..f377aa122 100644 --- a/cosmos_framework/data/generator/augmentors/reasoner/format_describe_anything.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/format_describe_anything.py @@ -178,6 +178,7 @@ def __init__( "random", ] = "random", urls_needs_timestamp: list = ["tl_plm_sav_20250714"], + stable_candidate_order: bool = False, ) -> None: """ Args: @@ -186,6 +187,7 @@ def __init__( self.input_key = input_key self.output_format = output_format self.urls_needs_timestamp = urls_needs_timestamp + self.stable_candidate_order: bool = stable_candidate_order def __call__(self, data_dict: Dict) -> Dict: url = data_dict["__url__"] @@ -243,8 +245,10 @@ def __call__(self, data_dict: Dict) -> Dict: if output_format in ["caption_one_object"]: assistant_message = random.sample(assistant_message, 1) elif output_format in ["location_and_caption_json_one_category", "location_and_caption_plain_one_category"]: - available_category = list( - set([assistant_message_i["category"] for assistant_message_i in assistant_message]) + candidates = [assistant_message_i["category"] for assistant_message_i in assistant_message] + # Offline replay must not depend on a worker's Python hash seed. + available_category = ( + list(dict.fromkeys(candidates)) if self.stable_candidate_order else list(set(candidates)) ) # sample one subject id category = random.choice(available_category) diff --git a/cosmos_framework/data/generator/augmentors/reasoner/prompt_format.py b/cosmos_framework/data/generator/augmentors/reasoner/prompt_format.py index 963502771..2eeea6a7f 100644 --- a/cosmos_framework/data/generator/augmentors/reasoner/prompt_format.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/prompt_format.py @@ -3,11 +3,13 @@ """Visual-Text Transformations or Augmentations.""" +import copy import random import re from typing import Any, Literal from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor +from cosmos_framework.utils.generator.structured_conversation import REASONING_PREFILL_KEY _THINK_RE = re.compile(r".*?\s*|.*\Z", re.DOTALL) @@ -18,18 +20,29 @@ def __init__( input_keys: list[str] = ["texts"], text_chat_order: Literal["text_end", "text_start", "random"] = "text_end", strip_thinking_prob: float = 0.0, + require_single_conversation: bool = False, + mask_inserted_reasoning_prefix: bool = False, ) -> None: """ Args: input_keys: List of input keys. text_chat_order: Order of text items in user messages. strip_thinking_prob: Per-sample probability of dropping assistant thinking traces. + require_single_conversation: Require a nested input to contain + exactly one conversation and unwrap it without random choice. + mask_inserted_reasoning_prefix: Exclude the inserted and newline + from loss. Opt in only to reproduce a preprocessing contract that + used this mask; otherwise preserve the processor-specific mask. """ if not 0.0 <= strip_thinking_prob <= 1.0: raise ValueError(f"strip_thinking_prob must be in [0, 1], got {strip_thinking_prob}") self.input_keys = input_keys self.text_chat_order = text_chat_order self.strip_thinking_prob = strip_thinking_prob + self.require_single_conversation = require_single_conversation + if not isinstance(mask_inserted_reasoning_prefix, bool): + raise TypeError("mask_inserted_reasoning_prefix must be a bool") + self.mask_inserted_reasoning_prefix: bool = mask_inserted_reasoning_prefix def __call__(self, data_dict: dict[str, Any]) -> dict[str, Any] | None: conversation_key = self.input_keys[0] @@ -45,7 +58,15 @@ def __call__(self, data_dict: dict[str, Any]) -> dict[str, Any] | None: # check if this is list of list of dict or list of dict if isinstance(list_of_conversation[0], list): - selected_index = random.sample(range(len(list_of_conversation)), 1)[0] + if self.require_single_conversation: + if len(list_of_conversation) != 1: + raise ValueError( + "Deterministic conversation-row preprocessing requires exactly one " + f"conversation, got {len(list_of_conversation)}." + ) + selected_index = 0 + else: + selected_index = random.sample(range(len(list_of_conversation)), 1)[0] selected_conversation = list_of_conversation[selected_index] conversation_metadata = data_dict.pop("_conversation_metadata", None) if conversation_metadata is not None: @@ -61,6 +82,7 @@ def __call__(self, data_dict: dict[str, Any]) -> dict[str, Any] | None: assert isinstance(selected_conversation, list) and isinstance(selected_conversation[0], dict), ( f"selected_conversation is not a list of dict: {selected_conversation}" ) + selected_conversation = copy.deepcopy(selected_conversation) # Normalize all string content to list format for message in selected_conversation: if "content" in message and isinstance(message["content"], str): @@ -96,7 +118,7 @@ def __call__(self, data_dict: dict[str, Any]) -> dict[str, Any] | None: has_media = any( isinstance(item, dict) and item.get("type") in ("image", "video", "audio") for item in content ) - if not has_text and not has_media: + if not has_text and not has_media and not message.get("tool_calls"): return None else: # Merge reasoning_content into assistant message content @@ -107,6 +129,8 @@ def __call__(self, data_dict: dict[str, Any]) -> dict[str, Any] | None: think_start = [{"type": "text", "text": "\n"}] think_end = [{"type": "text", "text": "\n\n\n"}] message["content"] = think_start + reasoning_items + think_end + message["content"] + if self.mask_inserted_reasoning_prefix: + message[REASONING_PREFILL_KEY] = True del message["reasoning_content"] data_dict["is_thinking_stripped"] = is_thinking_stripped diff --git a/cosmos_framework/data/generator/augmentors/reasoner/prompt_format_test.py b/cosmos_framework/data/generator/augmentors/reasoner/prompt_format_test.py index 9e89da57c..5abe9fa58 100644 --- a/cosmos_framework/data/generator/augmentors/reasoner/prompt_format_test.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/prompt_format_test.py @@ -1,9 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 +import copy + import pytest from cosmos_framework.data.generator.augmentors.reasoner.prompt_format import PromptFormat +from cosmos_framework.utils.generator.structured_conversation import REASONING_PREFILL_KEY pytestmark = [pytest.mark.L0, pytest.mark.CPU] @@ -73,3 +76,35 @@ def test_strip_thinking_drops_sample_without_assistant_supervision() -> None: def test_strip_thinking_rejects_invalid_probability(probability: float) -> None: with pytest.raises(ValueError, match=r"must be in \[0, 1\]"): PromptFormat(strip_thinking_prob=probability) + + +@pytest.mark.parametrize("mask_prefix", [False, True]) +def test_reasoning_prefix_supervision_is_explicit(mask_prefix: bool) -> None: + source = { + "texts": [ + {"role": "user", "content": "Question"}, + {"role": "assistant", "reasoning_content": "Reasoning", "content": "Answer", "loss": False}, + ] + } + original = copy.deepcopy(source) + source_messages = source["texts"] + result = PromptFormat(mask_inserted_reasoning_prefix=mask_prefix)(source) + assert result is not None + message = result["conversation"][1] + assert message.get(REASONING_PREFILL_KEY, False) is mask_prefix + assert message["loss"] is False + assert "".join(item["text"] for item in message["content"]) == "\nReasoning\n\n\nAnswer" + assert source_messages == original["texts"] + + +def test_reasoning_prefix_default_retains_processor_mask() -> None: + result = PromptFormat()( + { + "texts": [ + {"role": "user", "content": "Question"}, + {"role": "assistant", "reasoning_content": "Reasoning", "content": "Answer"}, + ] + } + ) + assert result is not None + assert REASONING_PREFILL_KEY not in result["conversation"][1] diff --git a/cosmos_framework/data/generator/augmentors/reasoner/sample_data_stats.py b/cosmos_framework/data/generator/augmentors/reasoner/sample_data_stats.py index 55c7f3b32..3f5a844c1 100644 --- a/cosmos_framework/data/generator/augmentors/reasoner/sample_data_stats.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/sample_data_stats.py @@ -11,6 +11,7 @@ import torch from cosmos_framework.data.generator.reasoner.video_decoder_qwen import pixels_to_token, token_to_pixels +from cosmos_framework.utils.generator.reasoner.constant import MEDIA_ROLES def _mean(values: list[float] | list[int]) -> float | None: @@ -57,7 +58,7 @@ def from_conversation(cls, conversation: list[dict[str, Any]]) -> "SampleDataSta num_videos = 0 for message in conversation: content = message.get("content") - if message.get("role") != "user" or not isinstance(content, list): + if message.get("role") not in MEDIA_ROLES or not isinstance(content, list): continue num_images += sum(item.get("type") == "image" for item in content) num_videos += sum(item.get("type") == "video" for item in content) diff --git a/cosmos_framework/data/generator/augmentors/reasoner/timestamp.py b/cosmos_framework/data/generator/augmentors/reasoner/timestamp.py index 986544f67..22e8c5b35 100644 --- a/cosmos_framework/data/generator/augmentors/reasoner/timestamp.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/timestamp.py @@ -35,9 +35,8 @@ class _TimestampProcessor(Protocol): def compute_timestamps(frame_index: int, fps: float, processor: _TimestampProcessor | None) -> float: - if processor is not None and ( - getattr(processor, "USES_TEMPORAL_PATCH_TIMESTAMPS", False) or "Qwen3" in processor.name - ): + use_patch_clock = processor is not None and (getattr(processor, "USES_TEMPORAL_PATCH_TIMESTAMPS", False) or "Qwen3" in processor.name) + if use_patch_clock: # Snapshot paths may be arbitrary hashes. Use the processor's temporal # contract, not its directory name or spatial merge size, for this clock. patch_size = getattr(processor, "temporal_patch_size", processor.merge_size) diff --git a/cosmos_framework/data/generator/augmentors/reasoner/timestamp_with_subject_tracking.py b/cosmos_framework/data/generator/augmentors/reasoner/timestamp_with_subject_tracking.py index 7496747e0..e8ce68ac0 100644 --- a/cosmos_framework/data/generator/augmentors/reasoner/timestamp_with_subject_tracking.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/timestamp_with_subject_tracking.py @@ -324,6 +324,7 @@ def __init__( ] = "random", urls_needs_timestamp: list = ["tl_plm_sav_20250714"], processor=None, + stable_candidate_order: bool = False, ) -> None: """ Args: @@ -332,6 +333,7 @@ def __init__( self.input_key = input_key self.output_format = output_format self.urls_needs_timestamp = urls_needs_timestamp + self.stable_candidate_order: bool = stable_candidate_order self.processor = processor def __call__(self, data_dict: Dict) -> Dict: @@ -406,8 +408,10 @@ def __call__(self, data_dict: Dict) -> Dict: ]: assistant_message = random.sample(assistant_message, 1) elif output_format in ["dense_video_caption_json_one_subject", "dense_video_caption_plain_one_subject"]: - available_subject_ids = list( - set([assistant_message_i["subject_id"] for assistant_message_i in assistant_message]) + candidates = [assistant_message_i["subject_id"] for assistant_message_i in assistant_message] + # Offline replay must not depend on a worker's Python hash seed. + available_subject_ids = ( + list(dict.fromkeys(candidates)) if self.stable_candidate_order else list(set(candidates)) ) # sample one subject id subject_id = random.choice(available_subject_ids) diff --git a/cosmos_framework/data/generator/augmentors/reasoner/tokenize_data.py b/cosmos_framework/data/generator/augmentors/reasoner/tokenize_data.py index 0d3d7c1ca..21de1df95 100644 --- a/cosmos_framework/data/generator/augmentors/reasoner/tokenize_data.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/tokenize_data.py @@ -9,6 +9,7 @@ import numpy as np import torch from PIL import Image +from qwen_vl_utils.vision_process import smart_resize from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log @@ -35,8 +36,9 @@ get_qwen_video_timestamps, splice_audio_segments_after_video_chunks, ) +from cosmos_framework.data.generator.processors.qwen3vl_processor import VIDEO_BUDGET_RESIZED_KEY from cosmos_framework.data.generator.processors.qwen3vl_processor import Qwen3VLProcessor as Processor -from cosmos_framework.utils.generator.reasoner.constant import IGNORE_INDEX, PROCESSOR_KEYS_TO_ADD +from cosmos_framework.utils.generator.reasoner.constant import IGNORE_INDEX, MEDIA_ROLES, PROCESSOR_KEYS_TO_ADD from cosmos_framework.utils.generator.source_video_timing import ( SOURCE_VIDEO_TIMING_KEY, require_source_pts_processor, @@ -76,6 +78,28 @@ def expand_video_frames_for_framewise( return expanded_videos, expanded_indices +def _select_video_frames(content: dict[str, Any], indices: list[int]) -> None: + """Select pixels and their timing metadata together, without mutating source records.""" + num_frames = len(content["video"]) + if "frames_indices" in content: + if len(content["frames_indices"]) != num_frames: + raise ValueError("Framewise metadata must have one source index per video frame") + content["frames_indices"] = [content["frames_indices"][index] for index in indices] + if VIDEO_METADATA_KEY in content: + metadata = validate_source_video_metadata(content[VIDEO_METADATA_KEY], num_frames) + metadata["frames_indices"] = [metadata["frames_indices"][index] for index in indices] + content[VIDEO_METADATA_KEY] = metadata + if SOURCE_VIDEO_TIMING_KEY in content: + timing = content[SOURCE_VIDEO_TIMING_KEY] + validate_source_video_timing(timing, num_frames) + selected_timing = dict(timing) + for key in ("frame_indices", "pts_seconds", "duration_seconds"): + selected_timing[key] = [timing[key][index] for index in indices] + selected_timing["origin_pts_seconds"] = selected_timing["pts_seconds"][0] + content[SOURCE_VIDEO_TIMING_KEY] = selected_timing + content["video"] = [content["video"][index] for index in indices] + + def maybe_subsample_frames(model_name_or_path, list_of_pil_image, max_video_token_length, processor): """ Why do we need to subsample frames? For model like eagle_er, it does not support smart downsampling in the processor. @@ -156,8 +180,13 @@ class TokenizeData(Augmentor): def __init__( self, processor: Optional[Processor] = None, + min_video_token_length: int = 16, max_video_token_length: int = 8192, + max_video_frames: int | None = None, max_image_token_length: int = 8192, + max_total_token_length: int | None = None, + adaptive_video_token_budget: bool = False, + adaptive_video_token_safety_margin: int = 128, custom_system_prompt: str | None = None, strip_original_system_prompt: bool = False, data_stats_writer_enabled: bool = False, @@ -172,6 +201,8 @@ def __init__( audio_layout: str = "separate_with_timestamps", video_timestamp_mode: str = "qwen_index", video_temporal_mode: VideoTemporalMode = "native", + add_vision_id: bool = False, + keep_raw_media: bool = True, ) -> None: """ Args: @@ -182,6 +213,13 @@ def __init__( data_stats_writer_enabled: Collect per-sample metadata for DataStatsWriterCallback. sound_und (bool): Opt in to audio preprocessing and audio-token registration. Disabled by default so existing text/vision tokenizers are unchanged. + add_vision_id (bool): Label media in samples with several images or several videos + ("Picture N: " / "Video N: "). Off by default: the processor wrapper used to drop this flag, so + checkpoints trained so far never saw the labels, and serving stacks do not add them by default. + It applies to every dataset this augmentor tokenizes, and datasets whose prompts already + interleave "Picture N" text would get duplicate labels, so it is not a recipe-wide setting. + keep_raw_media (bool): Also emit the debug-only "raw_image" / "raw_video" uint8 tensors. Disable for + long-context runs, where they duplicate every decoded frame in the dataloader output. """ validate_video_timestamp_mode(video_timestamp_mode) self.video_timestamp_mode: str = video_timestamp_mode @@ -192,11 +230,43 @@ def __init__( # Create the tokenizer self.text_only = text_only self.processor = processor # Expecting a ImageTextTokenizer + self.min_video_token_length = min_video_token_length self.max_video_token_length = max_video_token_length + self.max_video_frames = max_video_frames self.max_image_token_length = max_image_token_length + self.max_total_token_length = max_total_token_length + if not isinstance(adaptive_video_token_budget, bool): + raise TypeError( + f"adaptive_video_token_budget must be a bool, got {type(adaptive_video_token_budget).__name__}" + ) + if min_video_token_length <= 0: + raise ValueError(f"min_video_token_length must be positive, got {min_video_token_length}") + if max_video_token_length < min_video_token_length: + raise ValueError( + "max_video_token_length must be greater than or equal to min_video_token_length, " + f"got min={min_video_token_length}, max={max_video_token_length}" + ) + if max_video_frames is not None and max_video_frames <= 0: + raise ValueError(f"max_video_frames must be positive, got {max_video_frames}") + if adaptive_video_token_safety_margin < 0: + raise ValueError( + f"adaptive_video_token_safety_margin must be non-negative, got {adaptive_video_token_safety_margin}" + ) + if adaptive_video_token_budget and ( + max_total_token_length is None or max_total_token_length <= adaptive_video_token_safety_margin + ): + raise ValueError( + "adaptive_video_token_budget requires max_total_token_length greater than " + f"adaptive_video_token_safety_margin, got max_total_token_length={max_total_token_length}, " + f"margin={adaptive_video_token_safety_margin}" + ) + self.adaptive_video_token_budget = adaptive_video_token_budget + self.adaptive_video_token_safety_margin = adaptive_video_token_safety_margin self.custom_system_prompt = custom_system_prompt self.strip_original_system_prompt = strip_original_system_prompt self.collect_data_stats = data_stats_writer_enabled + self.add_vision_id = add_vision_id + self.keep_raw_media = keep_raw_media self.video_temporal_mode: VideoTemporalMode = video_temporal_mode self.effective_temporal_patch_size: int = get_effective_temporal_patch_size( getattr(self.processor, "temporal_patch_size", 2), video_temporal_mode @@ -226,6 +296,134 @@ def __init__( audio_end_token=audio_end_token, ) + def _video_token_counts(self, tokenizer_output: dict[str, torch.Tensor]) -> list[int]: + """Count post-merge tokens for each video represented by video_grid_thw.""" + video_grid_thw = tokenizer_output.get("video_grid_thw") + if video_grid_thw is None: + return [] + grid = torch.as_tensor(video_grid_thw, dtype=torch.int64) + if grid.numel() == 0: + return [] + if grid.ndim == 1: + grid = grid.reshape(1, -1) + unmerged_tokens = grid.prod(dim=-1) + merge_area = int(self.processor.merge_size) ** 2 + if bool(torch.any(unmerged_tokens % merge_area != 0)): + raise ValueError( + f"video_grid_thw contains per-video counts {unmerged_tokens.tolist()}, not all divisible by " + f"merge_size**2={merge_area}" + ) + return [int(value) for value in (unmerged_tokens // merge_area).tolist()] + + def _count_video_tokens(self, tokenizer_output: dict[str, torch.Tensor]) -> int: + """Count aggregate post-merge video tokens represented by video_grid_thw.""" + return sum(self._video_token_counts(tokenizer_output)) + + def _apply_video_token_budget( + self, + conversation: list[dict[str, Any]], + video_token_budget: int, + *, + allow_frame_subsampling: bool, + ) -> bool: + """Apply one aggregate video budget across all videos in a conversation.""" + video_contents = [ + content + for message in conversation + if isinstance(message.get("content"), list) + for content in message["content"] + if content.get("type") == "video" + ] + if not video_contents: + return False + + patch_size = int(self.processor.patch_size) + temporal_patch_size = int(self.processor.temporal_patch_size) + merge_size = int(self.processor.merge_size) + total_pixel_budget = token_to_pixels( + video_token_budget, + patch_size=patch_size, + temporal_patch_size=temporal_patch_size, + merge_size=merge_size, + ) + pixel_budget_per_video = total_pixel_budget // len(video_contents) + min_pixels_per_frame = self._minimum_video_pixels_per_frame() + + for content in video_contents: + frames = content["video"] + if allow_frame_subsampling: + max_frames = pixel_budget_per_video // min_pixels_per_frame + max_frames = max_frames // temporal_patch_size * temporal_patch_size + if max_frames < temporal_patch_size: + return False + if len(frames) > max_frames: + if self.video_temporal_mode == "framewise": + # Keep the repeated tubelet for each selected source frame intact. + groups = np.linspace( + 0, len(frames) // temporal_patch_size - 1, max_frames // temporal_patch_size + ).astype(int) + sample_idx = [ + int(group) * temporal_patch_size + offset + for group in groups + for offset in range(temporal_patch_size) + ] + else: + sample_idx = np.linspace(0, len(frames) - 1, max_frames).astype(int).tolist() + _select_video_frames(content, sample_idx) + frames = content["video"] + + max_pixels_per_frame = pixel_budget_per_video // len(frames) + if max_pixels_per_frame < min_pixels_per_frame: + return False + content["max_pixels"] = max_pixels_per_frame + if self.processor.use_smart_resize: + resized_frames = self._resize_video_frames(frames, max_pixels_per_frame) + if resized_frames is None: + return False + content["video"] = resized_frames + return True + + def _minimum_video_pixels_per_frame(self) -> int: + """Return the model-legal video floor without changing the decoder's training-quality floor.""" + factor = int(self.processor.patch_size) * int(self.processor.merge_size) + wrapped_processor = getattr(self.processor, "processor", None) + video_processor = getattr(wrapped_processor, "video_processor", None) + video_size = getattr(video_processor, "size", {}) + shortest_edge = video_size.get("shortest_edge") if isinstance(video_size, dict) else None + if shortest_edge is not None: + return max(int(shortest_edge), factor**2) + min_height_width = int(getattr(self.processor, "min_height_width", factor)) + return max(min_height_width**2, factor**2) + + def _resize_video_frames( + self, + frames: list[Image.Image], + max_pixels_per_frame: int, + ) -> list[Image.Image] | None: + """Downscale frames so processors that ignore content max_pixels still honor the budget.""" + patch_size = int(self.processor.patch_size) + merge_size = int(self.processor.merge_size) + factor = patch_size * merge_size + min_pixels_per_frame = self._minimum_video_pixels_per_frame() + if max_pixels_per_frame < min_pixels_per_frame: + return None + + resized_frames: list[Image.Image] = [] + for frame in frames: + width, height = frame.size + if width * height <= max_pixels_per_frame: + resized_frames.append(frame) + continue + resized_height, resized_width = smart_resize( + height, + width, + factor=factor, + min_pixels=min_pixels_per_frame, + max_pixels=max_pixels_per_frame, + ) + resized_frames.append(frame.resize((resized_width, resized_height), resample=Image.Resampling.BICUBIC)) + return resized_frames + def __call__(self, data_dict: Dict) -> Dict: r"""Tokenize a dialog and pad the sequence. @@ -282,7 +480,7 @@ def __call__(self, data_dict: Dict) -> Dict: raise ValueError( f"message is not a dict: {message} | conversation: {conversation} | data_dict: {data_dict} | __url__: {data_dict['__url__'].root}, {data_dict['__url__'].path}" ) - if message["role"] == "user" and isinstance(message["content"], list): + if message["role"] in MEDIA_ROLES and isinstance(message["content"], list): total_images += len([content for content in message["content"] if content["type"] == "image"]) total_videos += len([content for content in message["content"] if content["type"] == "video"]) total_audios += len([content for content in message["content"] if content["type"] == "audio"]) @@ -307,7 +505,7 @@ def __call__(self, data_dict: Dict) -> Dict: return None for message in conversation: - if message["role"] != "user" or not isinstance(message["content"], list): + if message["role"] not in MEDIA_ROLES or not isinstance(message["content"], list): continue for content in message["content"]: if content["type"] != "audio": @@ -357,8 +555,8 @@ def __call__(self, data_dict: Dict) -> Dict: for message in conversation: video_timestamps_by_key: dict[str, list[float]] = {} active_video_key: str | None = None - # for user message, we insert the media - if message["role"] == "user" and isinstance( + # for user and tool messages, we insert the media + if message["role"] in MEDIA_ROLES and isinstance( message["content"], list ): # Otherwise it's text and content is a string images_content_idx_full = [ @@ -393,7 +591,7 @@ def __call__(self, data_dict: Dict) -> Dict: min_pixels_per_image = self.processor.processor.image_processor.size["shortest_edge"] if max_pixels_per_image < min_pixels_per_image: log.critical( - f"max_pixels_per_image: {max_pixels_per_image} < min_pixels_per_image: {min_pixels_per_image} | self.max_video_token_length = {self.max_video_token_length} is not enough for total_images: {total_images}, as the default min_pixels is {min_pixels_per_image} | Either increase max_video_token_length or include max_pixels in the content or reduce min_pixels" + f"max_pixels_per_image: {max_pixels_per_image} < min_pixels_per_image: {min_pixels_per_image} | self.max_image_token_length = {self.max_image_token_length} is not enough for total_images: {total_images}, as the default min_pixels is {min_pixels_per_image} | Either increase max_image_token_length or include max_pixels in the content or reduce min_pixels" ) return None @@ -414,9 +612,10 @@ def __call__(self, data_dict: Dict) -> Dict: image = data_dict["media"][content["image"]] content["image"] = image content["max_pixels"] = max_pixels_per_image - raw_image = np.asarray(image.convert("RGB")) # [H,W,3] - raw_images.append(torch.from_numpy(raw_image).permute(2, 0, 1)[:, None]) # [3,1,H,W] - native_height, native_width = raw_image.shape[:2] + native_width, native_height = image.size + if self.keep_raw_media: + raw_image = np.asarray(image.convert("RGB")) # [H,W,3] + raw_images.append(torch.from_numpy(raw_image).permute(2, 0, 1)[:, None]) # [3,1,H,W] if data_stats is not None: data_stats.add_image(native_height, native_width, self.processor) @@ -424,7 +623,9 @@ def __call__(self, data_dict: Dict) -> Dict: # as tokenization will NOT upsample the video, we can use a larger value here at the cost of multiple video having 1.5x token length max_total_pixels = token_to_pixels( self.max_video_token_length * 1.5, + patch_size=self.processor.patch_size, temporal_patch_size=self.effective_temporal_patch_size, + merge_size=self.processor.merge_size, ) media_key = content["video"] # Add each video to the content list @@ -467,11 +668,42 @@ def __call__(self, data_dict: Dict) -> Dict: self.processor, self.effective_temporal_patch_size, ) - # this is because videos are decoded to be around "max_video_token_length" tokens - videos = maybe_subsample_frames( - self.processor.name, videos, self.max_video_token_length, self.processor + content["video"] = videos + if video_metadata is not None: + content[VIDEO_METADATA_KEY] = video_metadata + if source_timing is not None: + content[SOURCE_VIDEO_TIMING_KEY] = source_timing + if self.video_temporal_mode == "framewise": + content["frames_indices"] = ( + video_metadata["frames_indices"] + if video_metadata is not None + else video_media.get("source_frames_indices", list(range(len(videos)))) + ) + selected_indices = maybe_subsample_frames( + self.processor.name, list(range(len(videos))), self.max_video_token_length, self.processor ) + if self.max_video_frames is not None and len(selected_indices) > self.max_video_frames: + # Framewise frames are expanded to whole tubelets later, so each source frame counts. + temporal_patch_size = self.effective_temporal_patch_size + aligned_max_video_frames = ( + self.max_video_frames // temporal_patch_size * temporal_patch_size + ) + if aligned_max_video_frames < temporal_patch_size: + log.critical( + "[TokenizerDataError]max_video_frames must allow at least " + f"{temporal_patch_size} temporally aligned frames, " + f"got {self.max_video_frames}. url: {url}", + rank0_only=False, + ) + return None + sample_idx = np.linspace(0, len(selected_indices) - 1, aligned_max_video_frames).astype(int) + selected_indices = [selected_indices[index] for index in sample_idx] + if len(selected_indices) != len(videos): + _select_video_frames(content, selected_indices) + videos = content["video"] + video_metadata = content.get(VIDEO_METADATA_KEY) + source_timing = content.get(SOURCE_VIDEO_TIMING_KEY) if len(videos) == 0: log.info(f"[TokenizerDataError]video {media_key} has no decoded frames. url: {url}") return None @@ -482,14 +714,23 @@ def __call__(self, data_dict: Dict) -> Dict: if video_metadata is not None: # Any later temporal selection must update source indices with the pixels. content[VIDEO_METADATA_KEY] = validate_source_video_metadata(video_metadata, len(videos)) + content[VIDEO_BUDGET_RESIZED_KEY] = video_media.get(VIDEO_BUDGET_RESIZED_KEY, False) max_pixels_per_image = max_total_pixels // total_videos // len(videos) + # Explicit adaptive budgeting owns this pixel change. Ordinary + # loaders keep the original decoded frames and processor path. + if self.adaptive_video_token_budget and self.processor.use_smart_resize: + resized_videos = self._resize_video_frames(videos, int(max_pixels_per_image)) + if resized_videos is None: + log.critical( + "[TokenizerDataError]video pixel budget is below the processor's minimum " + f"frame geometry: max_pixels_per_image={max_pixels_per_image}. url: {url}", + rank0_only=False, + ) + return None + videos = resized_videos if self.video_temporal_mode == "framewise": - source_indices = ( - video_metadata["frames_indices"] - if video_metadata is not None - else video_media.get("source_frames_indices", list(range(len(videos)))) - ) + source_indices = content["frames_indices"] if len(source_indices) != len(videos): log.critical( f"[TokenizerDataError]video frame metadata length {len(source_indices)} does not " @@ -513,7 +754,6 @@ def __call__(self, data_dict: Dict) -> Dict: else: content["video"] = videos content["fps"] = fps - content["max_pixels"] = max_pixels_per_image if message_has_audio: video_timestamps_by_key[media_key] = get_qwen_video_timestamps( num_frames=len(videos), @@ -525,10 +765,11 @@ def __call__(self, data_dict: Dict) -> Dict: if self.audio_layout == "interleaved_av": audio_segment_lengths_by_video.append(None) - raw_video_frames = np.stack( - [np.asarray(frame.convert("RGB")) for frame in videos], axis=0 - ) # [T,H,W,3] - raw_videos.append(torch.from_numpy(raw_video_frames).permute(3, 0, 1, 2)) # [3,T,H,W] + if self.keep_raw_media: + raw_video_frames = np.stack( + [np.asarray(frame.convert("RGB")) for frame in videos], axis=0 + ) # [T,H,W,3] + raw_videos.append(torch.from_numpy(raw_video_frames).permute(3, 0, 1, 2)) # [3,T,H,W] elif content["type"] == "audio": assert audio_outputs is not None assert self.audio_special_tokens is not None @@ -620,6 +861,10 @@ def __call__(self, data_dict: Dict) -> Dict: if len(raw_videos) > 0: data_dict["raw_video"] = raw_videos # each: [3,T,H,W] + # Retain the tool menu even when source system prose is stripped. + if conversation and conversation[0].get("tools") is not None: + processor_kwargs["tools"] = conversation[0]["tools"] + if self.strip_original_system_prompt: conversation = [msg for msg in conversation if msg.get("role") != "system"] data_dict["conversation"] = conversation @@ -634,10 +879,7 @@ def __call__(self, data_dict: Dict) -> Dict: ) return None - if total_images > 1 or total_videos > 1: - add_vision_id = True - else: - add_vision_id = False + add_vision_id = self.add_vision_id and (total_images > 1 or total_videos > 1) try: conversation = convert_all_images_to_rgb(conversation) @@ -649,7 +891,7 @@ def __call__(self, data_dict: Dict) -> Dict: ) return None - try: + def tokenize_conversation() -> dict[str, torch.Tensor]: tokenizer_output = self.processor.apply_chat_template( conversation, tokenize=True, @@ -676,6 +918,118 @@ def __call__(self, data_dict: Dict) -> Dict: tokenizer_output["attention_mask"] = tokenizer_output["attention_mask"].new_ones( spliced_input_ids.shape ) + return tokenizer_output + + adaptive_video_budget_applied = False + try: + tokenizer_output = tokenize_conversation() + if self.adaptive_video_token_budget and total_videos > 0: + assert self.max_total_token_length is not None + initial_video_token_counts = self._video_token_counts(tokenizer_output) + if len(initial_video_token_counts) != total_videos: + raise ValueError(f"Expected {total_videos} video grids, got {len(initial_video_token_counts)}.") + minimum_retained_video_tokens = tuple( + min(self.min_video_token_length, token_count) for token_count in initial_video_token_counts + ) + minimum_video_tokens = sum(minimum_retained_video_tokens) + current_video_budget = self._count_video_tokens(tokenizer_output) + for _ in range(3): + total_token_length = int(tokenizer_output["input_ids"].numel()) + if total_token_length <= self.max_total_token_length: + break + + actual_video_tokens = self._count_video_tokens(tokenizer_output) + non_video_tokens = total_token_length - actual_video_tokens + hard_available_video_tokens = self.max_total_token_length - non_video_tokens + available_video_tokens = ( + self.max_total_token_length - self.adaptive_video_token_safety_margin - non_video_tokens + ) + if hard_available_video_tokens < minimum_video_tokens: + log.critical( + "[TokenizerDataError]text and non-video tokens leave too little room for video: " + f"total={total_token_length}, non_video={non_video_tokens}, " + f"hard_available_video={hard_available_video_tokens}, " + f"minimum_video={minimum_video_tokens}, max_total={self.max_total_token_length}. " + f"url: {url}", + rank0_only=False, + ) + return None + preferred_video_budget = max( + available_video_tokens, + minimum_video_tokens + self.adaptive_video_token_safety_margin, + ) + next_video_budget = min( + self.max_video_token_length, + hard_available_video_tokens, + preferred_video_budget, + actual_video_tokens - 1, + ) + if next_video_budget >= actual_video_tokens: + overflow = total_token_length - self.max_total_token_length + next_video_budget = actual_video_tokens - max( + 1, + overflow + self.adaptive_video_token_safety_margin, + ) + if next_video_budget < minimum_video_tokens: + log.critical( + "[TokenizerDataError]text and non-video tokens leave too little room for video: " + f"total={total_token_length}, non_video={non_video_tokens}, " + f"available_video={available_video_tokens}, minimum_video={minimum_video_tokens}, " + f"max_total={self.max_total_token_length}. url: {url}", + rank0_only=False, + ) + return None + if not self._apply_video_token_budget( + conversation, + next_video_budget, + allow_frame_subsampling=total_audios == 0, + ): + log.critical( + "[TokenizerDataError]adaptive video budget is below the processor's minimum " + f"frame geometry: budget={next_video_budget}. url: {url}", + rank0_only=False, + ) + return None + adaptive_video_budget_applied = True + current_video_budget = next_video_budget + tokenizer_output = tokenize_conversation() + else: + total_token_length = int(tokenizer_output["input_ids"].numel()) + if total_token_length > self.max_total_token_length: + log.critical( + "[TokenizerDataError]adaptive video tokenization did not fit after 3 retries: " + f"total={total_token_length}, max_total={self.max_total_token_length}, " + f"last_video_budget={current_video_budget}. url: {url}", + rank0_only=False, + ) + return None + + if int(tokenizer_output["input_ids"].numel()) > self.max_total_token_length: + log.critical( + "[TokenizerDataError]adaptive video tokenization exceeded the exact sequence limit: " + f"total={int(tokenizer_output['input_ids'].numel())}, " + f"max_total={self.max_total_token_length}. url: {url}", + rank0_only=False, + ) + return None + if adaptive_video_budget_applied: + final_video_token_counts = self._video_token_counts(tokenizer_output) + if len(final_video_token_counts) != len(minimum_retained_video_tokens) or any( + actual < minimum + for actual, minimum in zip( + final_video_token_counts, + minimum_retained_video_tokens, + strict=True, + ) + ): + log.critical( + "[TokenizerDataError]adaptive resize fell below the per-video quality floor after " + "patch-grid quantization: " + f"actual={final_video_token_counts}, minimum={minimum_retained_video_tokens}. " + f"url: {url}", + rank0_only=False, + ) + return None except Exception as e: if self.video_timestamp_mode == "source_pts": raise ValueError("source_pts failed while rendering aligned video timestamps") from e @@ -683,6 +1037,21 @@ def __call__(self, data_dict: Dict) -> Dict: f"Error in tokenizer_output: {e} | conversation: {conversation} | __url__: {url} | data_dict: {data_dict.keys()}" ) return None + + if adaptive_video_budget_applied and self.keep_raw_media: + raw_videos = [] + for message in conversation: + if not isinstance(message.get("content"), list): + continue + for content in message["content"]: + if content.get("type") != "video": + continue + raw_video_frames = np.stack( + [np.asarray(frame.convert("RGB")) for frame in content["video"]], + axis=0, + ) # [T,H,W,3] + raw_videos.append(torch.from_numpy(raw_video_frames).permute(3, 0, 1, 2)) # [3,T,H,W] + data_dict["raw_video"] = raw_videos input_ids = tokenizer_output["input_ids"] if audio_outputs is not None: assert self.audio_special_tokens is not None @@ -697,19 +1066,22 @@ def __call__(self, data_dict: Dict) -> Dict: ) return None if "image_grid_thw" in tokenizer_output and "raw_image" in data_dict: + patch_size = self.processor.patch_size resized_raw_images: list[torch.Tensor] = [] for raw_image, image_grid_thw in zip(data_dict["raw_image"], tokenizer_output["image_grid_thw"]): - # image_grid_thw: [t,h,w] + # image_grid_thw: [t,h,w] in patches _, h, w = image_grid_thw raw_image = torch.nn.functional.interpolate( - raw_image, size=(int(h) * 14, int(w) * 14), mode="bilinear", align_corners=False - ) # [3,1,h*14,w*14] + raw_image, size=(int(h) * patch_size, int(w) * patch_size), mode="bilinear", align_corners=False + ) # [3,1,h*P,w*P] resized_raw_images.append(raw_image) - data_dict["raw_image"] = resized_raw_images # each: [3,1,h*14,w*14] + data_dict["raw_image"] = resized_raw_images # each: [3,1,h*P,w*P] try: # token_mask: True for tokens to compute loss on; False for tokens to ignore token_mask = self.processor.add_assistant_tokens_mask(input_ids) + if hasattr(self.processor, "apply_supervision_overrides"): + token_mask = self.processor.apply_supervision_overrides(input_ids, token_mask, conversation) except Exception as e: log.critical( f"Error in add_assistant_tokens_mask: {e} | conversation: {conversation} | __url__: {url} | data_dict: {data_dict.keys()}" diff --git a/cosmos_framework/data/generator/augmentors/reasoner/tokenize_data_test.py b/cosmos_framework/data/generator/augmentors/reasoner/tokenize_data_test.py index b7ba5c195..12aa8d98d 100644 --- a/cosmos_framework/data/generator/augmentors/reasoner/tokenize_data_test.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/tokenize_data_test.py @@ -6,6 +6,7 @@ import wave from copy import deepcopy from types import SimpleNamespace +from typing import Any import numpy as np import pytest @@ -15,12 +16,19 @@ from cosmos_framework.data.generator.augmentors.reasoner.bytes_to_media import BytesToMedia from cosmos_framework.data.generator.augmentors.reasoner.filter_output_key import FilterOutputKey from cosmos_framework.data.generator.augmentors.reasoner.filter_seq_length import FilterSeqLength -from cosmos_framework.data.generator.augmentors.reasoner.tokenize_data import TokenizeData +from cosmos_framework.data.generator.augmentors.reasoner.tokenize_data import TokenizeData, _select_video_frames from cosmos_framework.data.generator.processors.audio_utils import ( AUDIO_END_TOKEN, AUDIO_PAD_TOKEN, AUDIO_START_TOKEN, ) +from cosmos_framework.data.generator.processors.qwen3vl_processor import VIDEO_BUDGET_RESIZED_KEY +from cosmos_framework.utils.generator.source_video_timing import ( + SOURCE_VIDEO_TIMING_KEY, + build_source_video_timing, + validate_source_video_timing, +) +from cosmos_framework.utils.generator.video_source_metadata import VIDEO_METADATA_KEY, validate_source_video_metadata class _FakeTokenizer: @@ -67,9 +75,11 @@ def __init__(self, *, name: str = "fake-vlm", extra_vocabulary: dict[str, int] | self.name = name self.tokenizer = _FakeTokenizer(extra_vocabulary) self.last_conversation = None + self.last_kwargs = None def apply_chat_template(self, conversation, **kwargs) -> dict[str, torch.Tensor]: self.last_conversation = deepcopy(conversation) + self.last_kwargs = kwargs input_ids: list[int] = [] video_grid_thw: list[list[int]] = [] special_tokens = (AUDIO_START_TOKEN, AUDIO_PAD_TOKEN, AUDIO_END_TOKEN) @@ -196,8 +206,10 @@ def test_bytes_to_media_decodes_audio_bytes_to_mono_16khz() -> None: @pytest.mark.L0 @pytest.mark.CPU +@pytest.mark.parametrize("use_start_frame_end_frame", [False, True]) def test_bytes_to_media_extracts_video_audio_only_when_requested( monkeypatch: pytest.MonkeyPatch, + use_start_frame_end_frame: bool, ) -> None: waveform = torch.tensor([0.25, -0.5], dtype=torch.float32) monkeypatch.setattr( @@ -211,15 +223,61 @@ def test_bytes_to_media_extracts_video_audio_only_when_requested( lambda *args, **kwargs: waveform, ) - video_only = BytesToMedia(is_input_pickle_byptes=False)( - {"media": {"clip.mp4": b"encoded-video"}}, + data = { + "media": {"clip.mp4": b"encoded-video", "clip.wav": b"encoded-audio"}, + "start_frame": 0, + "end_frame": 2, + } + video_only = BytesToMedia( + is_input_pickle_byptes=False, extract_audio=False, use_start_frame_end_frame=use_start_frame_end_frame + )( + deepcopy(data), ) - audio_visual = BytesToMedia(is_input_pickle_byptes=False, extract_audio=True)( - {"media": {"clip.mp4": b"encoded-video"}}, + audio_visual = BytesToMedia( + is_input_pickle_byptes=False, extract_audio=True, use_start_frame_end_frame=use_start_frame_end_frame + )( + deepcopy(data), ) assert "audio" not in video_only["media"]["clip.mp4"] assert audio_visual["media"]["clip.mp4"]["audio"] is waveform + assert video_only["media"]["clip.wav"] is waveform + assert audio_visual["media"]["clip.wav"] is waveform + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("audio_layout", ["separate_no_timestamps", "separate_with_timestamps", "interleaved_av"]) +def test_video_without_soundtrack_keeps_tokens_when_audio_understanding_is_enabled(audio_layout: str) -> None: + data = { + "__key__": "video-without-soundtrack", + "__url__": SimpleNamespace(root="root", path="path"), + "conversation": [ + {"role": "user", "content": [{"type": "video", "video": "clip.mp4"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "answer"}]}, + ], + "media": {"clip.mp4": {"videos": [Image.new("RGB", (4, 4)) for _ in range(2)], "fps": 4.0}}, + } + # The test double implements the processor methods used here without loading model assets. + baseline = TokenizeData( + processor=_FakeVLMProcessor(), # pyright: ignore[reportArgumentType] + sound_und=False, + max_video_token_length=16, + )(deepcopy(data)) + audio_processor = _FakeAudioProcessor() + output = TokenizeData( + processor=_FakeVLMProcessor(), # pyright: ignore[reportArgumentType] + sound_und=True, + audio_processor=audio_processor, + audio_layout=audio_layout, + max_video_token_length=16, + )(deepcopy(data)) + + assert baseline is not None and output is not None + for key in ("input_ids", "attention_mask", "token_mask", "labels", "video_grid_thw", "pixel_values_videos"): + assert torch.equal(output[key], baseline[key]), key + assert "audio_features" not in output + assert audio_processor.last_audios is None @pytest.mark.L0 @@ -648,3 +706,225 @@ def _build_audio_processor(selected_audio_encoder_type: str) -> _FakeAudioProces assert constructed == [audio_encoder_type] assert tokenize.audio_processor is audio_processor + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize( + "framewise,cap,expected_indices,expected_pixels", + [ + (False, 2, [10, 40], [0, 3]), + (True, 2, [10, 10, 40, 40], [0, 0, 3, 3]), + (False, 3, [10, 40], [0, 3]), + (True, 3, [10, 10, 20, 20, 40, 40], [0, 0, 1, 1, 3, 3]), + ], +) +def test_video_frame_cap_keeps_source_indices_aligned( + framewise: bool, cap: int, expected_indices: list[int], expected_pixels: list[int] +) -> None: + processor = _FakeVLMProcessor() + tokenize = TokenizeData( + processor=processor, + max_video_frames=cap, + video_temporal_mode="framewise" if framewise else "native", + ) + metadata = {"fps": 10.0, "total_num_frames": 100, "frames_indices": [10, 20, 30, 40]} + data = { + "__key__": "frame-cap", + "__url__": SimpleNamespace(root="root", path="path"), + "conversation": [ + {"role": "user", "content": [{"type": "video", "video": "clip"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "answer"}]}, + ], + "media": { + "clip": { + "videos": [Image.new("RGB", (28, 28), (index, 0, 0)) for index in range(4)], + "fps": 1.0, + VIDEO_METADATA_KEY: metadata, + } + }, + } + assert tokenize(data) is not None + assert processor.last_conversation is not None + video = processor.last_conversation[0]["content"][0] + assert video[VIDEO_METADATA_KEY]["frames_indices"] == expected_indices + assert [frame.getpixel((0, 0))[0] for frame in video["video"]] == expected_pixels + assert validate_source_video_metadata(video[VIDEO_METADATA_KEY], len(video["video"]))["fps"] == 10.0 + assert metadata["frames_indices"] == [10, 20, 30, 40] + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_tokenize_data_resolves_and_budgets_images_returned_by_tools() -> None: + vlm_processor = _FakeVLMProcessor() + data = { + "__key__": "agentic-trajectory", + "__url__": SimpleNamespace(root="root", path="path"), + "media": {"screenshot": Image.new("RGB", (64, 48)), "crop": Image.new("RGB", (16, 16))}, + "conversation": [ + {"role": "user", "content": [{"type": "image", "image": "screenshot"}, {"type": "text", "text": "before"}]}, + { + "role": "assistant", + "content": "middle", + "tool_calls": [{"type": "function", "function": {"name": "crop", "arguments": {}}}], + }, + {"role": "tool", "content": [{"type": "image", "image": "crop"}]}, + {"role": "assistant", "content": "answer"}, + ], + } + + output = TokenizeData(processor=vlm_processor, max_image_token_length=1024)(data) + + assert output is not None + user_image = vlm_processor.last_conversation[0]["content"][0] + tool_image = vlm_processor.last_conversation[2]["content"][0] + assert isinstance(tool_image["image"], Image.Image) + # Both images share the image budget: 1,024 tokens of 28x28 merged patches, halved. + assert user_image["max_pixels"] == tool_image["max_pixels"] == 1024 * 28 * 28 // 2 + + +def _two_video_sample() -> dict: + frames = [Image.new("RGB", (32, 32)) for _ in range(2)] + return { + "__key__": "two-videos", + "__url__": SimpleNamespace(root="root", path="path"), + "media": { + "decoded": {"videos": list(frames), "fps": 2.0, VIDEO_BUDGET_RESIZED_KEY: True}, + "other": {"videos": list(frames), "fps": 2.0}, + }, + "conversation": [ + { + "role": "user", + "content": [ + {"type": "video", "video": "decoded"}, + {"type": "video", "video": "other"}, + {"type": "text", "text": "before"}, + ], + }, + {"role": "assistant", "content": "answer"}, + ], + } + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("framewise", [False, True]) +def test_adaptive_budget_preserves_metadata_and_framewise_tubelets(framewise: bool) -> None: + tokenize = TokenizeData(processor=_FakeVLMProcessor(), video_temporal_mode="framewise" if framewise else "native") + pixel_ids = [0, 0, 1, 1, 2, 2, 3, 3] if framewise else list(range(8)) + source_indices = [10 * (index + 1) for index in pixel_ids] + content: dict[str, Any] = { + "type": "video", + "video": [Image.new("RGB", (28, 28), (index, 0, 0)) for index in pixel_ids], + VIDEO_METADATA_KEY: {"fps": 10.0, "total_num_frames": 100, "frames_indices": source_indices}, + } + if framewise: + content["frames_indices"] = source_indices.copy() + conversation = [{"role": "user", "content": [content]}] + assert tokenize._apply_video_token_budget(conversation, 2, allow_frame_subsampling=True) + expected_pixels = [0, 0, 3, 3] if framewise else [0, 2, 4, 7] + expected_indices = [10 * (index + 1) for index in expected_pixels] + assert [frame.getpixel((0, 0))[0] for frame in content["video"]] == expected_pixels + assert content[VIDEO_METADATA_KEY]["frames_indices"] == expected_indices + validate_source_video_metadata(content[VIDEO_METADATA_KEY], len(content["video"])) + if framewise: + assert content["frames_indices"] == expected_indices + assert source_indices == [10 * (index + 1) for index in pixel_ids] + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_tokenize_data_passes_on_decoder_resizing_and_can_skip_raw_media() -> None: + vlm_processor = _FakeVLMProcessor() + + output = TokenizeData(processor=vlm_processor, keep_raw_media=False)(_two_video_sample()) + + assert output is not None + videos = vlm_processor.last_conversation[0]["content"][:2] + # Copied per video; the real Qwen3VLProcessor then rejects such a mix (see qwen3vl_processor_test). + assert [video[VIDEO_BUDGET_RESIZED_KEY] for video in videos] == [True, False] + assert "raw_video" not in output + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_temporal_selection_keeps_source_pts_aligned() -> None: + timing = build_source_video_timing( + [10, 20, 30, 40], + torch.tensor([3.0, 4.0, 5.0, 6.5], dtype=torch.float64), + torch.tensor([0.1, 0.2, 0.3, 0.4], dtype=torch.float64), + "a" * 64, + ) + original = deepcopy(timing) + content: dict[str, Any] = {"video": [0, 1, 2, 3], SOURCE_VIDEO_TIMING_KEY: timing} + _select_video_frames(content, [1, 3]) + selected = content[SOURCE_VIDEO_TIMING_KEY] + assert content["video"] == [1, 3] + assert selected["frame_indices"] == [20, 40] + assert selected["duration_seconds"] == [0.2, 0.4] + assert validate_source_video_timing(selected, 2) == [0.0, 2.5] + assert timing == original + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("add_vision_id", [False, True]) +def test_tokenize_data_labels_media_only_when_enabled(add_vision_id: bool) -> None: + vlm_processor = _FakeVLMProcessor() + + TokenizeData(processor=vlm_processor, add_vision_id=add_vision_id)(_two_video_sample()) + + assert vlm_processor.last_kwargs["add_vision_id"] is add_vision_id + + +class _GridProcessor(_FakeVLMProcessor): + """One input id per merged video token (grid [frame pairs, 2, 2]) and one per text item.""" + + def apply_chat_template(self, conversation, **kwargs) -> dict[str, torch.Tensor]: + self.last_conversation = deepcopy(conversation) + ids: list[int] = [] + grids: list[list[int]] = [] + for message in conversation: + for content in message["content"] if isinstance(message["content"], list) else []: + if content["type"] == "video": + pairs = (len(content["video"]) + 1) // 2 + grids.append([pairs, 2, 2]) + ids += [self.video_token_id] * pairs + elif content["type"] == "text": + ids.append(1) + input_ids = torch.tensor(ids, dtype=torch.long) + output = {"input_ids": input_ids, "attention_mask": torch.ones_like(input_ids, dtype=torch.bool)} + if grids: + output["video_grid_thw"] = torch.tensor(grids) + return output + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("keep_raw_media", [True, False]) +def test_adaptive_video_budget_honors_keep_raw_media(keep_raw_media: bool) -> None: + frames = [Image.new("RGB", (28, 28), (index, 0, 0)) for index in range(16)] + sample = { + "__key__": "adaptive-raw-media", + "__url__": SimpleNamespace(root="root", path="path"), + "media": {"clip": {"videos": frames, "fps": 2.0}}, + "conversation": [ + {"role": "user", "content": [{"type": "video", "video": "clip"}, {"type": "text", "text": "before"}]}, + {"role": "assistant", "content": "answer"}, + ], + } + processor = _GridProcessor() + tokenize = TokenizeData( + processor=processor, + keep_raw_media=keep_raw_media, + adaptive_video_token_budget=True, + max_total_token_length=6, + adaptive_video_token_safety_margin=0, + min_video_token_length=1, + ) + + output = tokenize(sample) + + assert output is not None + assert len(processor.last_conversation[0]["content"][0]["video"]) < len(frames) # the budget was applied + assert ("raw_video" in output) is keep_raw_media diff --git a/cosmos_framework/data/generator/augmentors/target_first_frame_test.py b/cosmos_framework/data/generator/augmentors/target_first_frame_test.py new file mode 100644 index 000000000..cad64ca20 --- /dev/null +++ b/cosmos_framework/data/generator/augmentors/target_first_frame_test.py @@ -0,0 +1,48 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import json + +import pytest +import torch + +from cosmos_framework.data.generator.augmentors.video_editing_transform import PairedVideoEditingToTrainingFormat +from cosmos_framework.data.generator.sequence_packing.packers import resolve_item_condition_frames +from cosmos_framework.utils.generator.torchcodec_video import VideoMetadata + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + + +@pytest.mark.parametrize("probability", [0.0, 1.0]) +def test_clean_prefix_changes_only_final_target_mask(monkeypatch: pytest.MonkeyPatch, probability: float) -> None: + transform = PairedVideoEditingToTrainingFormat(target_first_frame_probability=probability) + source = torch.zeros(3, 9, 16, 32) # [C,T,H,W] + target = torch.arange(source.numel()).reshape_as(source).float() # [C,T,H,W] + original = target.clone() + metadata = VideoMetadata(num_frames=9, average_fps=16.0, height=16, width=32) + monkeypatch.setattr(transform, "_decode_pair", lambda *_: (source, target, metadata, 32, 16)) + conversation = json.dumps( + [ + [ + { + "role": "user", + "content": [{"type": "video", "video": "source"}, {"type": "text", "text": "Add a cup."}], + }, + {"role": "assistant", "content": [{"type": "video", "video": "target"}]}, + ] + ] + ) + result = transform({"texts": conversation, "media": {"source": b"source", "target": b"target"}}) + assert result is not None + assert result["video"][0] is source + assert result["video"][1] is target + torch.testing.assert_close(target, original) + indexes = result["sequence_plan"].condition_frame_indexes_vision + assert indexes == ([0] if probability else []) + assert resolve_item_condition_frames(indexes, item_idx=0, num_items=2, latent_t=3) == [0, 1, 2] + assert resolve_item_condition_frames(indexes, item_idx=1, num_items=2, latent_t=3) == indexes + + +@pytest.mark.parametrize("probability", [-1, 2, float("nan")]) +def test_invalid_probability_rejected(probability: float) -> None: + with pytest.raises(ValueError, match="probability"): + PairedVideoEditingToTrainingFormat(target_first_frame_probability=probability) diff --git a/cosmos_framework/data/generator/augmentors/text_transforms_for_video.py b/cosmos_framework/data/generator/augmentors/text_transforms_for_video.py index 046992c6c..27bc65ffb 100644 --- a/cosmos_framework/data/generator/augmentors/text_transforms_for_video.py +++ b/cosmos_framework/data/generator/augmentors/text_transforms_for_video.py @@ -8,8 +8,8 @@ from collections.abc import Mapping, Sequence from typing import Optional -from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.data.imaginaire.webdataset.augmentors.v3_text_transforms import pad_and_resize +from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log diff --git a/cosmos_framework/data/generator/augmentors/torchcodec_callers_test.py b/cosmos_framework/data/generator/augmentors/torchcodec_callers_test.py index 7912de9dc..e1a2dc067 100644 --- a/cosmos_framework/data/generator/augmentors/torchcodec_callers_test.py +++ b/cosmos_framework/data/generator/augmentors/torchcodec_callers_test.py @@ -190,3 +190,5 @@ def test_multiview_augmentations_defer_normalization_to_model() -> None: augmentations = module.make_auto_augmentations({}) assert not any("normalize" in type(augmentor).__name__.lower() for augmentor in augmentations) + + diff --git a/cosmos_framework/data/generator/augmentors/video_editing_transform.py b/cosmos_framework/data/generator/augmentors/video_editing_transform.py index c0604fbfa..570ba2204 100644 --- a/cosmos_framework/data/generator/augmentors/video_editing_transform.py +++ b/cosmos_framework/data/generator/augmentors/video_editing_transform.py @@ -229,7 +229,9 @@ class PairedVideoEditingToTrainingFormat(Augmentor): Source and target frames are sampled on the same ``target_fps`` timeline and resized to the source aspect ratio. Their encoded durations may differ by at most one target-frame interval plus ``duration_ratio_tolerance`` of the longer - clip. Reference images are emitted between source and target as independent + clip. Explicit ``bernini_77f_source_sampling`` reproduces the generator input + timeline for native-rate 77-80f sources paired with a 77f target. Reference + images are emitted between source and target as independent one-frame vision items. """ @@ -244,10 +246,12 @@ def __init__( max_pixels: int = 512 * 512, padding_divisor: int = 16, duration_ratio_tolerance: float = 0.0, + bernini_77f_source_sampling: bool = False, aspect_ratio_tolerance: float = 0.05, mean: float = 0.5, std: float = 0.5, num_decode_threads: int = 1, + target_first_frame_probability: float = 0.0, max_reference_images: int = 4, min_reference_image_side: int = 128, args: dict | None = None, @@ -261,6 +265,12 @@ def __init__( raise ValueError("min_reference_image_side must be positive") if not 0 <= duration_ratio_tolerance < 1 or not 0 <= aspect_ratio_tolerance < 1: raise ValueError("alignment tolerances must be in [0, 1)") + if not 0 <= target_first_frame_probability <= 1: + raise ValueError("target_first_frame_probability must be in [0, 1]") + self.target_first_frame_probability = target_first_frame_probability + self.bernini_77f_source_sampling = bernini_77f_source_sampling + if bernini_77f_source_sampling and (max_num_frames != 77 or target_fps != 16.0): + raise ValueError("Bernini source sampling requires 77 target frames at 16 FPS") self.media_key = media_key self.conversation_key = conversation_key self.target_fps = target_fps @@ -327,7 +337,14 @@ def _validate_pair(self, source: VideoMetadata, target: VideoMetadata) -> tuple[ # Editing pairs should be temporally aligned. Permit one target-frame interval # for container/FPS metadata rounding, plus an explicitly configured ratio. duration_tolerance = 1.0 / self.target_fps + self.duration_ratio_tolerance * longer_duration - if duration_difference > duration_tolerance: + # Bernini 77f generation samples the native-rate source using 4N+1 + # linspace indices before truncation. Opt in only for its known 77-80f + # cohort; longer targets and arbitrary FPS/duration mismatches stay invalid. + # The source-only tail is bounded; neither clip is stretched in time. + bernini_pair = self.bernini_77f_source_sampling and self._is_bernini_77f_pair(source, target) + if self.bernini_77f_source_sampling and not bernini_pair: + raise ValueError("bernini_77f_pair_requires_77_80_source_77_target_at_16fps") + if duration_difference > duration_tolerance and not bernini_pair: raise ValueError(f"duration_mismatch:{source_duration:.3f}:{target_duration:.3f}") source_aspect = source.width / source.height @@ -349,6 +366,23 @@ def _validate_pair(self, source: VideoMetadata, target: VideoMetadata) -> tuple[ ) return common_duration, width, height + @staticmethod + def _is_bernini_77f_pair(source: VideoMetadata, target: VideoMetadata) -> bool: + return ( + 77 <= source.num_frames <= 80 + and target.num_frames == 77 + and math.isclose(source.average_fps, 16.0, rel_tol=0.0, abs_tol=1e-6) + and math.isclose(target.average_fps, 16.0, rel_tol=0.0, abs_tol=1e-6) + ) + + @staticmethod + def _bernini_source_frame_indices(source_frames: int) -> list[int]: + # Match e6c2cf11 smart_video_nframes(fps=16, frame_factor=4, + # max_frames=77, add_one=True): round the complete 4N+1 linspace FIRST, + # then truncate to 77. For 80f this repeats one middle source frame. + sampled_frames = 1 + 4 * (source_frames // 4) + return [int(index) for index in torch.linspace(0, source_frames - 1, sampled_frames).round().tolist()][:77] + def _decode_pair( self, source_bytes: bytes, target_bytes: bytes ) -> tuple[torch.Tensor, torch.Tensor, VideoMetadata, int, int]: @@ -356,9 +390,21 @@ def _decode_pair( target_metadata = probe_video(target_bytes, include_dimensions=True, num_threads=self.num_decode_threads) common_duration, width, height = self._validate_pair(source_metadata, target_metadata) available_frames = max(1, math.floor(common_duration * self.target_fps)) - num_frames = valid_video_frame_count(available_frames, self.max_num_frames) - source_indices = aligned_frame_indices(source_metadata, self.target_fps, num_frames) - target_indices = aligned_frame_indices(target_metadata, self.target_fps, num_frames) + # The validated Bernini cohort always contains exactly 77 target frames. + # Near-16 FPS metadata must not floor its duration down to 76 (then 73). + num_frames = ( + 77 if self.bernini_77f_source_sampling else valid_video_frame_count(available_frames, self.max_num_frames) + ) + source_indices = ( + self._bernini_source_frame_indices(source_metadata.num_frames) + if self.bernini_77f_source_sampling + else aligned_frame_indices(source_metadata, self.target_fps, num_frames) + ) + target_indices = ( + list(range(77)) + if self.bernini_77f_source_sampling + else aligned_frame_indices(target_metadata, self.target_fps, num_frames) + ) source_tchw, _ = decode_frames_tchw_uint8( # [T,C,Hs,Ws] source_bytes, source_indices, num_threads=self.num_decode_threads ) @@ -398,6 +444,11 @@ def __call__(self, data_dict: dict) -> dict | None: references = [self._decode_reference_image(media[key]) for key in reference_keys] reference_tensors = [tensor for tensor, _, _ in references] num_frames = source_cthw.shape[1] + if self.target_first_frame_probability and num_frames < 5: + raise ValueError("target first-frame conditioning requires generated frames after the prefix") + condition_target_first_frame = ( + self.target_first_frame_probability > 0 and random.random() < self.target_first_frame_probability + ) data_dict["video"] = [source_cthw, *reference_tensors, target_cthw] data_dict["ai_caption"] = instruction data_dict["selected_caption_type"] = "editing_instruction" @@ -415,7 +466,9 @@ def __call__(self, data_dict: dict) -> dict | None: data_dict["sequence_plan"] = SequencePlan( has_text=True, has_vision=True, - condition_frame_indexes_vision=[], + # In multi-item packing this applies to the final target only. + # The causal VAE encodes its first frame independently of future frames. + condition_frame_indexes_vision=[0] if condition_target_first_frame else [], # Source and target videos are frame-aligned. Reference images remain # independent because their one-frame temporal grid differs from the videos. share_vision_temporal_positions=not reference_tensors, diff --git a/cosmos_framework/data/generator/augmentors/video_editing_transform_test.py b/cosmos_framework/data/generator/augmentors/video_editing_transform_test.py index d5b191a8f..d3f8d325b 100644 --- a/cosmos_framework/data/generator/augmentors/video_editing_transform_test.py +++ b/cosmos_framework/data/generator/augmentors/video_editing_transform_test.py @@ -18,6 +18,8 @@ ) from cosmos_framework.utils.generator.torchcodec_video import VideoMetadata +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + def _conversation(instruction: str = "Remove the car", reference_keys: tuple[str, ...] = ()) -> str: return json.dumps( @@ -241,3 +243,80 @@ def test_pair_validation_rejects_duration_difference_over_one_target_frame() -> def test_parse_rejects_empty_instruction() -> None: with pytest.raises(ValueError, match="instruction is empty"): parse_video_editing_conversation(_conversation("")) + + +@pytest.mark.parametrize( + ("source_frames", "target_frames", "fps", "allowed"), + [ + (80, 77, 16.0, True), + (77, 77, 16.0, True), + (81, 77, 16.0, False), + (77, 80, 16.0, False), + (80, 77, 15.0, False), + (78, 75, 16.0, False), + ], +) +def test_bernini_source_tail_is_bounded_and_native_rate( + source_frames: int, target_frames: int, fps: float, allowed: bool +) -> None: + transform = PairedVideoEditingToTrainingFormat(max_num_frames=77, bernini_77f_source_sampling=True) + source = VideoMetadata(num_frames=source_frames, average_fps=fps, height=64, width=96) + target = VideoMetadata(num_frames=target_frames, average_fps=fps, height=64, width=96) + if allowed: + duration, _, _ = transform._validate_pair(source, target) + assert duration == pytest.approx(77 / 16) + assert aligned_frame_indices(source, 16.0, 77) == list(range(77)) + assert aligned_frame_indices(target, 16.0, 77) == list(range(77)) + else: + with pytest.raises(ValueError, match="bernini_77f_pair"): + transform._validate_pair(source, target) + + +def test_source_tail_remains_strict_without_explicit_opt_in() -> None: + source = VideoMetadata(num_frames=80, average_fps=16.0, height=64, width=96) + target = VideoMetadata(num_frames=77, average_fps=16.0, height=64, width=96) + with pytest.raises(ValueError, match="duration_mismatch"): + PairedVideoEditingToTrainingFormat(max_num_frames=77)._validate_pair(source, target) + + +@pytest.mark.parametrize(("maximum", "fps"), [(93, 16.0), (77, 15.0)]) +def test_bernini_sampling_requires_its_native_target(maximum: int, fps: float) -> None: + with pytest.raises(ValueError, match="Bernini source sampling"): + PairedVideoEditingToTrainingFormat(max_num_frames=maximum, target_fps=fps, bernini_77f_source_sampling=True) + + +def test_bernini_indices_match_generation_before_truncation() -> None: + indices = PairedVideoEditingToTrainingFormat._bernini_source_frame_indices(80) + assert len(indices) == 77 + assert indices == torch.linspace(0, 79, 81).round().long().tolist()[:77] + assert indices != list(range(77)) + assert indices[0] == 0 + assert indices[-1] == 75 + assert len(set(indices)) == 76 + assert PairedVideoEditingToTrainingFormat._bernini_source_frame_indices(77) == list(range(77)) + + +@pytest.mark.parametrize("fps", [16.0, 16.0000005, 15.9999995]) +@pytest.mark.parametrize("source_frames", [77, 80]) +def test_bernini_decode_keeps_both_videos_at_77_frames_despite_fps_rounding( + monkeypatch: pytest.MonkeyPatch, fps: float, source_frames: int +) -> None: + from cosmos_framework.data.generator.augmentors import video_editing_transform as module + + transform = PairedVideoEditingToTrainingFormat(max_num_frames=77, bernini_77f_source_sampling=True) + source_metadata = VideoMetadata(num_frames=source_frames, average_fps=fps, height=64, width=96) + target_metadata = VideoMetadata(num_frames=77, average_fps=fps, height=64, width=96) + monkeypatch.setattr( + module, "probe_video", lambda payload, **_: source_metadata if payload == b"source" else target_metadata + ) + selected: dict[bytes, list[int]] = {} + + def decode(payload: bytes, indices: list[int], **kwargs: object) -> tuple[torch.Tensor, VideoMetadata]: + selected[payload] = indices + return torch.zeros((len(indices), 3, 64, 96), dtype=torch.uint8), target_metadata + + monkeypatch.setattr(module, "decode_frames_tchw_uint8", decode) + source, target, _, _, _ = transform._decode_pair(b"source", b"target") + assert source.shape[1] == target.shape[1] == 77 + assert selected[b"target"] == list(range(77)) + assert selected[b"source"] == transform._bernini_source_frame_indices(source_frames) diff --git a/cosmos_framework/data/generator/indexed_transfer_length_loader.py b/cosmos_framework/data/generator/indexed_transfer_length_loader.py index a5eef16b5..ec543f688 100644 --- a/cosmos_framework/data/generator/indexed_transfer_length_loader.py +++ b/cosmos_framework/data/generator/indexed_transfer_length_loader.py @@ -53,7 +53,7 @@ def __init__( *, index_creds_path: str | Path | None = None, index_cache_dir: str | Path | None = None, - caption_key: str = "canonical_caption_structured", + caption_key: str = "caption", conditioning_config: dict[int, float] | None = None, resolution: str = "480", seed: int = 42, diff --git a/cosmos_framework/data/generator/joint_dataloader.py b/cosmos_framework/data/generator/joint_dataloader.py index d016c14d4..2c04b053d 100644 --- a/cosmos_framework/data/generator/joint_dataloader.py +++ b/cosmos_framework/data/generator/joint_dataloader.py @@ -5,10 +5,12 @@ import multiprocessing import queue import threading +import time from collections import deque -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field +from itertools import chain from typing import Any, ClassVar, Dict, Union import numpy as np @@ -20,6 +22,7 @@ from cosmos_framework.utils import log from cosmos_framework.utils.generator.cost_model.budget import IterationTimeBudget, IterationTimeBudgetConfig from cosmos_framework.data.generator.drop_sample_contract import DROP_SAMPLE_KEY, DROP_SAMPLE_REASON_KEY +from cosmos_framework.data.generator.sample_token_counter import SampleTokenCounter from cosmos_framework.data.generator.token_mix_control import TokenMixControlConfig, TokenMixController from cosmos_framework.data.generator.worker_timing import ( _TIMING_KEYS as _TIMING_KEYS, @@ -27,13 +30,6 @@ from cosmos_framework.data.generator.worker_timing import ( _aggregate_worker_timing as _aggregate_worker_timing, ) -from cosmos_framework.model.generator.tokenizers.uniae.frame_math import ( - get_uniae_chunk_frames, - get_uniae_latent_num_frames, - normalize_uniae_chunk_frames, -) -from cosmos_framework.utils.generator.data_utils import read_positive_int_metadata -from cosmos_framework.utils.generator.spatial_patch import normalize_spatial_patch_hw _BATCH_TIMING_KEYS = { "_worker_batch_time", @@ -260,7 +256,7 @@ class _PackingMetrics: """Per-batch packing statistics collected during the packing loop. Also serves as the single source of truth for packing-related metric names - via ``STATS_SPEC``, which the dataloading monitor callback consumes to + via ``STATS_SPEC``, which the packing stats monitor callback consumes to drive accumulation and logging. """ @@ -436,7 +432,7 @@ def close(self) -> None: self._thread.join(timeout=min(1.0, self._timeout_s)) -class JointDataLoader(webdataset.WebLoader): +class JointDataLoader(SampleTokenCounter, webdataset.WebLoader): r""" A joint dataloader that supports loading both images and videos. """ @@ -461,6 +457,7 @@ def __init__( audio_sample_rate: int = 48000, prewarm: bool = True, prewarm_concurrency: int = 1, + prewarm_rank_stagger_s: float = 0.0, default_lookahead_limit: int = _DEFAULT_LOOKAHEAD_LIMIT, lookahead_limits: Dict[str, int] | None = None, uniae_chunk_frames: int | Mapping[str, int] | None = None, @@ -470,6 +467,9 @@ def __init__( forkserver_preload_modules: list[str] | None = None, lidar_patch_spatial_hw: int | tuple[int, int] | None = None, radar_patch_spatial_hw: int | tuple[int, int] | None = None, + enable_async_batch_building: bool = False, + async_batch_building_timeout_s: float = 1200.0, + distributed_row_index_cache: bool = True, ) -> None: """ Initialize the JointDataLoader with multiple datasets. @@ -506,6 +506,11 @@ def __init__( sound token count. prewarm_concurrency: Number of independent dataloaders to prewarm at once. Keep this bounded to avoid an object-store request burst. + prewarm_rank_stagger_s: Maximum deterministic delay across distributed + ranks before child iterator startup. This smooths the job-wide + object-store request ramp without changing stream order. + distributed_row_index_cache: Prepare image/video/multiview caches with disjoint + node scans and row-ID redistribution before starting child workers. default_lookahead_limit: Packing-loop look-ahead fallback for dataloaders not in ``lookahead_limits``. lookahead_limits: Optional ``{dataset_name: int}`` per-dataloader override. @@ -520,6 +525,9 @@ def __init__( ``iteration_time_target`` is ``None``, packs on tokens alone. forkserver_preload_modules: Optional modules imported once by the multiprocessing forkserver before it creates child workers. + enable_async_batch_building: Build one future packed batch on a background + thread for loaders using the shared synchronous packing iterator. + async_batch_building_timeout_s: Maximum time to wait for an async batch. Example: joint_loader = IterativeJointDataLoader( @@ -535,67 +543,43 @@ def __init__( } ) """ + if async_batch_building_timeout_s <= 0: + raise ValueError(f"async_batch_building_timeout_s must be positive, got {async_batch_building_timeout_s}.") + # Set this before initializing child iterators. + self.enable_async_batch_building: bool = enable_async_batch_building + # Async batch builder state + self.async_batch_building_timeout_s: float = float(async_batch_building_timeout_s) + self._async_batch_builder: _AsyncBatchBuilder | None = None self.dataloader_list, self.dataset_name_list, self.data_ratios = [], [], [] self.lookahead_limits: list[int] = [] - self.tokenizer_spatial_compression_factor = tokenizer_spatial_compression_factor - self.tokenizer_temporal_compression_factor = tokenizer_temporal_compression_factor - self.lidar_spatial_compression = ( - tuple(int(factor) for factor in lidar_spatial_compression) - if lidar_spatial_compression is not None - else None - ) - if self.lidar_spatial_compression is not None and ( - len(self.lidar_spatial_compression) != 2 or any(factor <= 0 for factor in self.lidar_spatial_compression) - ): - raise ValueError( - "lidar_spatial_compression must contain two positive factors " - f"(height, width), got {self.lidar_spatial_compression}" - ) - self.lidar_temporal_compression_factor = ( - int(lidar_temporal_compression_factor) if lidar_temporal_compression_factor is not None else None - ) - if self.lidar_temporal_compression_factor is not None and self.lidar_temporal_compression_factor <= 0: - raise ValueError( - f"lidar_temporal_compression_factor must be positive, got {self.lidar_temporal_compression_factor}" - ) - self.radar_spatial_compression = ( - tuple(int(factor) for factor in radar_spatial_compression) - if radar_spatial_compression is not None - else None - ) - if self.radar_spatial_compression is not None and ( - len(self.radar_spatial_compression) != 2 or any(factor <= 0 for factor in self.radar_spatial_compression) - ): - raise ValueError( - "radar_spatial_compression must contain two positive factors " - f"(height, width), got {self.radar_spatial_compression}" - ) - self.radar_temporal_compression_factor = ( - int(radar_temporal_compression_factor) if radar_temporal_compression_factor is not None else None - ) - if self.radar_temporal_compression_factor is not None and self.radar_temporal_compression_factor <= 0: - raise ValueError( - f"radar_temporal_compression_factor must be positive, got {self.radar_temporal_compression_factor}" - ) - self.patch_spatial = patch_spatial - self.lidar_patch_spatial_hw: tuple[int, int] = normalize_spatial_patch_hw( - patch_spatial if lidar_patch_spatial_hw is None else lidar_patch_spatial_hw - ) - self.radar_patch_spatial_hw: tuple[int, int] = normalize_spatial_patch_hw( - patch_spatial if radar_patch_spatial_hw is None else radar_patch_spatial_hw + SampleTokenCounter.__init__( + self, + tokenizer_spatial_compression_factor=tokenizer_spatial_compression_factor, + tokenizer_temporal_compression_factor=tokenizer_temporal_compression_factor, + patch_spatial=patch_spatial, + lidar_spatial_compression=lidar_spatial_compression, + lidar_temporal_compression_factor=lidar_temporal_compression_factor, + radar_spatial_compression=radar_spatial_compression, + radar_temporal_compression_factor=radar_temporal_compression_factor, + sound_latent_fps=sound_latent_fps, + audio_sample_rate=audio_sample_rate, + uniae_chunk_frames=uniae_chunk_frames, + uniae_pad_frames=uniae_pad_frames, + lidar_patch_spatial_hw=lidar_patch_spatial_hw, + radar_patch_spatial_hw=radar_patch_spatial_hw, ) self.max_sequence_length = max_sequence_length self.max_samples_per_batch = max_samples_per_batch - self.sound_latent_fps = sound_latent_fps - self.audio_sample_rate = audio_sample_rate if prewarm_concurrency < 1: raise ValueError(f"prewarm_concurrency must be at least 1, got {prewarm_concurrency}.") + if not math.isfinite(prewarm_rank_stagger_s) or prewarm_rank_stagger_s < 0: + raise ValueError(f"prewarm_rank_stagger_s must be finite and nonnegative, got {prewarm_rank_stagger_s}.") self.prewarm_concurrency = prewarm_concurrency + self.prewarm_rank_stagger_s: float = float(prewarm_rank_stagger_s) self.forkserver_preload_modules: list[str] = list(forkserver_preload_modules or []) self.default_lookahead_limit = int(default_lookahead_limit) - self.uniae_pad_frames = int(uniae_pad_frames) if uniae_pad_frames is not None else None - self.uniae_chunk_frames = self._normalize_uniae_chunk_frames(uniae_chunk_frames) self.prewarm: bool = bool(prewarm) + self.distributed_row_index_cache = distributed_row_index_cache self.lazy_initialize_child_iterators: bool = bool(lazy_initialize_child_iterators) self._child_iterators_initialized: bool = False self._drop_sample_log_count: int = 0 @@ -649,17 +633,40 @@ def __init__( else: self._initialize_child_iterators_once() + def _validate_async_child_startup(self) -> None: + """Check worker startup even when a synchronous child is wrapped by a hybrid loader.""" + # Keep PyTorch/WebDataset/Lance worker creation outside the packer thread. + # Workerless loaders can initialize lazily without fetching every stream first. + if not self.prewarm: + if any(getattr(loader, "num_workers", None) != 0 for loader in self.dataloader_list): + raise ValueError("Async batch building without prewarm requires num_workers=0 on every child loader.") + def _initialize_child_iterators_once(self) -> None: """Create child iterators and optionally prewarm them once.""" if self._child_iterators_initialized: return - if self.forkserver_preload_modules: - preload_modules = list(dict.fromkeys(["__main__", *self.forkserver_preload_modules])) + if self.enable_async_batch_building: + self._validate_async_child_startup() + preload_modules = list(self.forkserver_preload_modules) + for dataloader in self.dataloader_list: + preload_modules.extend(getattr(dataloader, "forkserver_preload_modules", [])) + if preload_modules: + preload_modules = list(dict.fromkeys(["__main__", *preload_modules])) multiprocessing.set_forkserver_preload(preload_modules) log.info( f"JointDataLoader: configured forkserver preload modules: {preload_modules}", rank0_only=False, ) + if self.prewarm and self.prewarm_rank_stagger_s > 0 and torch.distributed.is_initialized(): + rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + delay_s = self.prewarm_rank_stagger_s * rank / world_size + log.info( + f"JointDataLoader: delaying child iterator startup by {delay_s:.1f}s " + f"(rank={rank}/{world_size}, maximum={self.prewarm_rank_stagger_s:.1f}s).", + rank0_only=False, + ) + time.sleep(delay_s) log.info( f"JointDataLoader: initializing {len(self.dataloader_list)} child iterator(s).", rank0_only=False, @@ -691,86 +698,6 @@ def _log_drop_sample(self, sample: Mapping[str, Any], dataset_name: str) -> None rank0_only=False, ) - def _num_lidar_tokens(self, data_batch: Mapping[str, Any]) -> int: - """Cost the sample's LiDAR range clips with the LiDAR VAE's own compression. - - The LiDAR VAE does not compress time, so costing a sweep clip with the camera's 4x - would undercount it fourfold and silently over-pack the batch. - """ - clips = data_batch.get("lidar") - if not clips: - return 0 - if self.lidar_spatial_compression is None or self.lidar_temporal_compression_factor is None: - raise ValueError( - "This batch carries a LiDAR stream, but the loader has no LiDAR compression factors. " - "Set lidar_spatial_compression and lidar_temporal_compression_factor." - ) - spatial_h, spatial_w = self.lidar_spatial_compression - num_tokens = 0 - for clip in clips: - _, T, H, W = clip.shape - patch_h = math.ceil(H // spatial_h / self.lidar_patch_spatial_hw[0]) - patch_w = math.ceil(W // spatial_w / self.lidar_patch_spatial_hw[1]) - latent_t = 1 + (T - 1) // self.lidar_temporal_compression_factor - num_tokens += patch_h * patch_w * latent_t - return num_tokens - - def _num_radar_tokens(self, data_batch: Mapping[str, Any]) -> int: - """Cost the sample's radar BEV clips with the radar VAE's own compression. - - Same arithmetic as the sweeps over a square BEV grid rather than a range image: the - radar VAE does not compress time either, so one scan is one latent frame and the - camera's 4x would undercount a clip fourfold. Radar is the more expensive of the two - sensors per second of clip -- it cycles at ~20 Hz against LiDAR's 10 and each scan - carries more latents -- so an underpriced clip over-packs the batch by more here. - """ - clips = data_batch.get("radar") - if not clips: - return 0 - if self.radar_spatial_compression is None or self.radar_temporal_compression_factor is None: - raise ValueError( - "This batch carries a radar stream, but the loader has no radar compression factors. " - "Set radar_spatial_compression and radar_temporal_compression_factor." - ) - spatial_h, spatial_w = self.radar_spatial_compression - num_tokens = 0 - for clip in clips: - _, T, H, W = clip.shape - patch_h = math.ceil(H // spatial_h / self.radar_patch_spatial_hw[0]) - patch_w = math.ceil(W // spatial_w / self.radar_patch_spatial_hw[1]) - latent_t = 1 + (T - 1) // self.radar_temporal_compression_factor - num_tokens += patch_h * patch_w * latent_t - return num_tokens - - def _normalize_uniae_chunk_frames( - self, uniae_chunk_frames: int | Mapping[str, int] | None - ) -> int | dict[str, int] | None: - return normalize_uniae_chunk_frames( - uniae_chunk_frames, - pad_frames=self.uniae_pad_frames, - temporal_compression_factor=self.tokenizer_temporal_compression_factor, - temporal_divisibility_name="tokenizer_temporal_compression_factor", - ) - - def _get_uniae_chunk_frames(self, spatial_shape: tuple[int, int]) -> int: - assert self.uniae_chunk_frames is not None - return get_uniae_chunk_frames(self.uniae_chunk_frames, spatial_shape=spatial_shape) - - def _compute_vision_latent_t_shape(self, T: int, H: int, W: int) -> int: - if T < 1: - raise ValueError(f"Vision media must contain at least one frame, got {T}.") - if T == 1 or self.uniae_chunk_frames is None: - return 1 + (T - 1) // self.tokenizer_temporal_compression_factor - - assert self.uniae_pad_frames is not None - return get_uniae_latent_num_frames( - T, - self.uniae_chunk_frames, - pad_frames=self.uniae_pad_frames, - temporal_compression_factor=self.tokenizer_temporal_compression_factor, - spatial_shape=(H, W), - ) - def _prewarm_dataloader(self, index: int, name: str, dl_iter: Any) -> None: """Produce and buffer one batch from a single dataloader.""" import time @@ -855,155 +782,6 @@ def _prewarm_dataloaders(self) -> None: torch.distributed.barrier() log.info("Pre-warm: all ranks ready") - def _compute_num_tokens_per_sample(self, data_batch: dict) -> int: - """ - This function computes the number of tokens per sample in the data batch. - This includes text + vision generation tokens + action tokens. - - Args: - data_batch (dict): The data batch containing the text tokens. - - Returns: - int: The number of tokens per sample. - """ - und_tokens, gen_tokens = self._compute_token_split_per_sample(data_batch) - return und_tokens + gen_tokens - - def _compute_token_split_per_sample(self, data_batch: dict) -> tuple[int, int]: - """ - Split one sample's packed length into the understanding and generation pathways. - - The two towers of the MoT attend over the same packed sequence but cost - different amounts per token, so the cost model prices them separately. The - two returned counts always sum to ``_compute_num_tokens_per_sample``. - - The three structural markers (````, ````, ````) - are attributed to the stream they delimit rather than broken out, which keeps - that sum exact at the cost of a handful of tokens landing on one side or the - other. - - Args: - data_batch (dict): The data batch containing the text tokens. - - Returns: - tuple[int, int]: ``(understanding tokens, generation tokens)``. - """ - - # The token sequence we have is - # [] - # The spatial dimension of image tokens is compressed by - # vae spatial downsampling factor + pathification - # The temporal dimension of image tokens is compressed by - # vae temporal downsampling factor - # Action tokens have 1 token per time step (no spatial dimension) - - has_text_tokens = "text_token_ids" in data_batch - und_tokens = 0 - gen_tokens = 0 - if has_text_tokens: - text_token_ids = data_batch["text_token_ids"] - if isinstance(text_token_ids, list): - # Price every independently tokenized per-view caption. The packing MR - # will use these boundaries as separate UND segments; accounting for the - # wrapper per segment here prevents the batch builder from underpacking. - und_tokens = sum(int(tokens.shape[0]) + 1 for tokens in text_token_ids) - else: - num_text_tokens = text_token_ids.shape[1] - und_tokens = num_text_tokens + 1 - - # Vision part. A LiDAR-only sample has no camera stream, so the loop below prices - # nothing and the whole generation cost comes from the sweeps. - is_image_batch = "images" in data_batch - # Absent rather than empty for a LiDAR-only sample, and tested for explicitly: the loop - # below also accepts a bare tensor, and `or []` would ask that tensor for its truth - # value before reaching that branch. - input_images_or_videos = data_batch.get("images" if is_image_batch else "video") - if input_images_or_videos is None: - input_images_or_videos = [] - if "enable_per_camera_vae_encoding" in data_batch: - sample_n_views_values = read_positive_int_metadata(data_batch, "sample_n_views", expected_count=1) - frames_per_view_values = read_positive_int_metadata( - data_batch, - "num_video_frames_per_view", - expected_count=1, - ) - if sample_n_views_values is None or frames_per_view_values is None: - raise ValueError( - "sample_n_views and num_video_frames_per_view are required when per-camera VAE encoding is enabled." - ) - sample_n_views = sample_n_views_values[0] - frames_per_view = frames_per_view_values[0] - else: - sample_n_views = None - frames_per_view = None - - # iterate over all the media in the batch - for media in input_images_or_videos if isinstance(input_images_or_videos, list) else [input_images_or_videos]: - if is_image_batch: - _, H, W = media.shape - T = 1 - else: - _, T, H, W = media.shape - - latent_h_shape = H // self.tokenizer_spatial_compression_factor - latent_w_shape = W // self.tokenizer_spatial_compression_factor - patch_h_shape = math.ceil(latent_h_shape / self.patch_spatial) - patch_w_shape = math.ceil(latent_w_shape / self.patch_spatial) - if sample_n_views is not None and frames_per_view is not None and not is_image_batch: - # The multiview dataset resizes every camera to the same H/W and - # selects the same synchronized frame count before concatenating - # the camera-major clips along T. - expected_frames = sample_n_views * frames_per_view - if T != expected_frames: - raise ValueError( - "Multiview vision length must equal sample_n_views * num_video_frames_per_view: " - f"got T={T}, sample_n_views={sample_n_views}, num_video_frames_per_view={frames_per_view}." - ) - latent_t_shape = sample_n_views * self._compute_vision_latent_t_shape(frames_per_view, H, W) - else: - latent_t_shape = self._compute_vision_latent_t_shape(T, H, W) - - num_vision_tokens = patch_h_shape * patch_w_shape * latent_t_shape - if has_text_tokens: - num_vision_tokens += 2 - gen_tokens += num_vision_tokens - - # LiDAR part: its own VAE, hence its own compression factors. Charged to - # gen_tokens so the iteration-time cost model sees the sweeps. - gen_tokens += self._num_lidar_tokens(data_batch) - - # Radar part: a third VAE and a third pair of factors, charged the same way. A joint - # camera + radar sample carries no sweeps and a joint camera + LiDAR one no scans, so - # each of the two calls prices nothing on the other's recipe. - gen_tokens += self._num_radar_tokens(data_batch) - - # Action part: each action time step is 1 token. - # Action tensor shape is (T_action, D) per sample; stored as a single-element list. - if "action" in data_batch: - list_of_actions = data_batch["action"] - for action in list_of_actions: - # skip None actions - if action is None: - continue - num_action_tokens = action.shape[0] - gen_tokens += num_action_tokens - - # Sound part — estimate sound tokens from audio waveform length - if self.sound_latent_fps > 0 and "sound" in data_batch: - sound_data = data_batch["sound"] - if isinstance(sound_data, list) and len(sound_data) > 0: - first_sound = sound_data[0] - # Unwrap nested list if needed - if isinstance(first_sound, list): - first_sound = first_sound[0] - if first_sound is not None and isinstance(first_sound, torch.Tensor): - num_audio_samples = first_sound.shape[-1] - audio_duration = num_audio_samples / self.audio_sample_rate - num_sound_tokens = int(audio_duration * self.sound_latent_fps) - gen_tokens += num_sound_tokens - - return und_tokens, gen_tokens - def _compute_sample_cost(self, data_batch: dict) -> tuple[int, float]: """ Packed length and projected iteration-time cost of one sample. @@ -1156,11 +934,57 @@ def _get_next_sample(self, index_id: int) -> dict: return buffer.popleft() - def set_start_iteration(self, iteration: int): + def set_start_iteration(self, iteration: int) -> None: + """Set the batch sequence and request async work if child iterators already exist.""" + if self._async_batch_builder is not None: + raise RuntimeError("set_start_iteration must be called before async batch building starts.") self.global_id = iteration + if self.enable_async_batch_building and self._child_iterators_initialized: + self._start_async_batch_builder() - def __iter__(self): - raise NotImplementedError("__iter__ function is not implemented yet") + def __iter__(self) -> Iterator[dict[str, Any]]: + self._initialize_child_iterators_once() + if not self.enable_async_batch_building: + return self._iter_synchronous() + return self._iter_asynchronous() + + def _iter_synchronous(self) -> Iterator[dict[str, Any]]: + raise NotImplementedError("The subclass must implement its packing iterator.") + + def _start_async_batch_builder(self) -> _AsyncBatchBuilder: + """Create the async builder and request its first batch.""" + if self._async_batch_builder is not None: + raise RuntimeError("Async batch building has already started.") + self._initialize_child_iterators_once() + async_batch_builder = _AsyncBatchBuilder( + self._iter_synchronous(), + timeout_s=self.async_batch_building_timeout_s, + ) + self._async_batch_builder = async_batch_builder + async_batch_builder.request_batch(self.global_id) + return async_batch_builder + + def _iter_asynchronous(self) -> Iterator[dict[str, Any]]: + """Yield each completed batch while the async builder builds its replacement.""" + async_batch_builder = self._async_batch_builder or self._start_async_batch_builder() + + while True: + try: + batch = async_batch_builder.get_batch() + except StopIteration: + break + + # Request one replacement before yielding. The request queue contains only a + # sequence ID, and the result queue returns this exact batch object by reference. + async_batch_builder.request_batch(self.global_id) + yield batch + + def close(self) -> None: + """Permanently stop the optional async batch builder.""" + if self._async_batch_builder is not None: + # Keep the closed builder reference: its bounded join may return while the daemon + # is still blocked in data loading, so starting a replacement could race with it. + self._async_batch_builder.close() class IterativeJointDataLoader(JointDataLoader): @@ -1205,6 +1029,7 @@ def __init__( seed: int | None = 42, prewarm: bool = True, prewarm_concurrency: int = 1, + prewarm_rank_stagger_s: float = 0.0, default_lookahead_limit: int = JointDataLoader._DEFAULT_LOOKAHEAD_LIMIT, lookahead_limits: Dict[str, int] | None = None, uniae_chunk_frames: int | Mapping[str, int] | None = None, @@ -1212,15 +1037,16 @@ def __init__( enable_async_batch_building: bool = False, async_batch_building_timeout_s: float = 1200.0, packing_limits: dict[str, dict[str, int]] | None = None, + batch_collate_fn: Callable[[list[dict[str, Any]]], dict[str, Any]] | None = None, lazy_initialize_child_iterators: bool = False, iteration_time_budget: IterationTimeBudgetConfig | None = None, token_mix_control: TokenMixControlConfig | None = None, forkserver_preload_modules: list[str] | None = None, lidar_patch_spatial_hw: int | tuple[int, int] | None = None, radar_patch_spatial_hw: int | tuple[int, int] | None = None, + distributed_row_index_cache: bool = True, ) -> None: - if async_batch_building_timeout_s <= 0: - raise ValueError(f"async_batch_building_timeout_s must be positive, got {async_batch_building_timeout_s}.") + self.batch_collate_fn: Callable[[list[dict[str, Any]]], dict[str, Any]] | None = batch_collate_fn self.packing_limits: dict[str, dict[str, int]] = { name: dict(limits) for name, limits in (packing_limits or {}).items() } @@ -1232,9 +1058,6 @@ def __init__( raise ValueError(f"Invalid packing limits for {name}: {limits}") if any(isinstance(value, bool) or not isinstance(value, int) or value <= 0 for value in limits.values()): raise ValueError(f"Packing limits for {name} must be positive integers: {limits}") - # Set this before the base constructor initializes child iterators. - self.enable_async_batch_building: bool = enable_async_batch_building - # A time ceiling skews the token mix here exactly as it does on the random # loader, but the correction cannot be applied here: this loader seeds the draw # so every rank picks the same modality, and weights estimated from a rank's own @@ -1273,6 +1096,8 @@ def __init__( audio_sample_rate=audio_sample_rate, prewarm=prewarm, prewarm_concurrency=prewarm_concurrency, + prewarm_rank_stagger_s=prewarm_rank_stagger_s, + distributed_row_index_cache=distributed_row_index_cache, default_lookahead_limit=default_lookahead_limit, lookahead_limits=lookahead_limits, uniae_chunk_frames=uniae_chunk_frames, @@ -1282,32 +1107,17 @@ def __init__( forkserver_preload_modules=forkserver_preload_modules, lidar_patch_spatial_hw=lidar_patch_spatial_hw, radar_patch_spatial_hw=radar_patch_spatial_hw, + enable_async_batch_building=enable_async_batch_building, + async_batch_building_timeout_s=async_batch_building_timeout_s, ) self.seed = seed # Calculate probabilities for random sampling total_ratio = sum(self.data_ratios) self.data_probs = np.array([ratio / total_ratio for ratio in self.data_ratios]) - # Async batch builder state - self.async_batch_building_timeout_s: float = float(async_batch_building_timeout_s) - self._async_batch_builder: _AsyncBatchBuilder | None = None - - def _initialize_child_iterators_once(self) -> None: - # Keep PyTorch/WebDataset/Lance worker creation outside the packer thread. - # Workerless loaders can initialize lazily without fetching every stream first. - if self.enable_async_batch_building and not self.prewarm: - if any(getattr(loader, "num_workers", None) != 0 for loader in self.dataloader_list): - raise ValueError("Async batch building without prewarm requires num_workers=0 on every child loader.") - super()._initialize_child_iterators_once() - - def __iter__(self) -> Iterator[dict[str, Any]]: - self._initialize_child_iterators_once() - if not self.enable_async_batch_building: - return self._iter_synchronous() - return self._iter_asynchronous() def _iter_synchronous(self) -> Iterator[dict[str, Any]]: - """Run the existing packing path on the calling thread.""" + """Pack batches on the caller or the sole async producer thread.""" while True: if self.seed is not None: rng = np.random.RandomState(self.seed + self.global_id) @@ -1318,6 +1128,7 @@ def _iter_synchronous(self) -> Iterator[dict[str, Any]]: metrics = _PackingMetrics() output_batch = dict() + samples: list[dict[str, Any]] = [] skipped_samples = deque() lookahead_limit = self.lookahead_limits[index_id] lookahead_count = 0 @@ -1333,7 +1144,7 @@ def _iter_synchronous(self) -> Iterator[dict[str, Any]]: break # If we have started packing and tried lookahead_limit times to find a fitting sample but failed, stop. - if len(output_batch) > 0 and lookahead_count >= lookahead_limit: + if metrics.num_samples > 0 and lookahead_count >= lookahead_limit: break had_buffer = len(self.buffers[index_id]) > 0 @@ -1362,10 +1173,10 @@ def _iter_synchronous(self) -> Iterator[dict[str, Any]]: sample_seconds=sample_seconds, packed_tokens=metrics.current_sequence_length, packed_sample_seconds=metrics.sample_seconds, - batch_started=len(output_batch) > 0, + batch_started=metrics.num_samples > 0, max_sequence_length=token_limit, ): - if len(output_batch) == 0: + if metrics.num_samples == 0: # Only the token limit can reject an empty batch, so this sample does not # fit at any batch size: discard it and get the next sample. log.info( @@ -1387,16 +1198,23 @@ def _iter_synchronous(self) -> Iterator[dict[str, Any]]: metrics.sample_seconds += sample_seconds metrics.num_samples += 1 output["dataset_name"] = self.dataset_name_list[index_id] - self._update_output_batch(output_batch, output) + if self.batch_collate_fn is not None: + samples.append(output) + else: + self._update_output_batch(output_batch, output) # Add back skipped samples to the buffer for the next batch. # appendleft puts item at HEAD. So we insert S3, then S2, then S1. for sample in reversed(skipped_samples): self.buffers[index_id].appendleft(sample) - if len(output_batch) == 0: + if metrics.num_samples == 0: return + # Heterogeneous pools use the same aligned format as random-mixed + # loading; ordinary iterative recipes retain their existing format. + if self.batch_collate_fn is not None: + output_batch = self.batch_collate_fn(samples) metrics.attach_to( output_batch, buffer_size=len(self.buffers[index_id]), @@ -1405,49 +1223,6 @@ def _iter_synchronous(self) -> Iterator[dict[str, Any]]: self.global_id += 1 yield output_batch - def _start_async_batch_builder(self) -> _AsyncBatchBuilder: - """Create the async builder and request its first batch.""" - if self._async_batch_builder is not None: - raise RuntimeError("Async batch building has already started.") - self._initialize_child_iterators_once() - async_batch_builder = _AsyncBatchBuilder( - self._iter_synchronous(), - timeout_s=self.async_batch_building_timeout_s, - ) - self._async_batch_builder = async_batch_builder - async_batch_builder.request_batch(self.global_id) - return async_batch_builder - - def _iter_asynchronous(self) -> Iterator[dict[str, Any]]: - """Yield each completed batch while the async builder builds its replacement.""" - async_batch_builder = self._async_batch_builder or self._start_async_batch_builder() - - while True: - try: - batch = async_batch_builder.get_batch() - except StopIteration: - break - - # Request one replacement before yielding. The request queue contains only a - # sequence ID, and the result queue returns this exact batch object by reference. - async_batch_builder.request_batch(self.global_id) - yield batch - - def set_start_iteration(self, iteration: int) -> None: - """Set the batch sequence and request async work if child iterators already exist.""" - if self._async_batch_builder is not None: - raise RuntimeError("set_start_iteration must be called before async batch building starts.") - self.global_id = iteration - if self.enable_async_batch_building and self._child_iterators_initialized: - self._start_async_batch_builder() - - def close(self) -> None: - """Permanently stop the optional async batch builder.""" - if self._async_batch_builder is not None: - # Keep the closed builder reference: its bounded join may return while the daemon - # is still blocked in data loading, so starting a replacement could race with it. - self._async_batch_builder.close() - def _get_dataloader_index(self, data_id: float) -> int: """Maps global id to the corresponding dataloader index based on ratio.""" for i, r in enumerate(self.data_ratios): @@ -1771,6 +1546,7 @@ def __init__( uniae_pad_frames: int | None = None, iteration_time_budget: IterationTimeBudgetConfig | None = None, token_mix_control: TokenMixControlConfig | None = None, + distributed_row_index_cache: bool = True, ): super().__init__( dataloaders, @@ -1781,6 +1557,7 @@ def __init__( max_samples_per_batch, sound_latent_fps=sound_latent_fps, audio_sample_rate=audio_sample_rate, + distributed_row_index_cache=distributed_row_index_cache, default_lookahead_limit=default_lookahead_limit, lookahead_limits=lookahead_limits, uniae_chunk_frames=uniae_chunk_frames, @@ -1901,11 +1678,16 @@ def __iter__(self): @dataclass(frozen=True) class SampleTokenMixConfig: - """Estimate source token lengths per fresh sample, including deferred draws. + """Estimate source token lengths per sample accepted into a packed batch. Ratios become target token shares when enabled. The estimate is updated once - per fresh, packable candidate; rechecking a deferred sample never counts it a - second time. Both loaders use the shared ``TokenMixController`` estimator. + per accepted sample; deferred candidates do not update the estimate, warmup, + or feedback until accepted. Cumulative accepted token shares also adjust + subsequent draw weights toward the target. This feedback is best effort and + does not restrict fallback. + ``half_life_samples`` controls only the token-length estimate; feedback counts + tokens since the loader was created. Both loaders use the shared + ``TokenMixController`` estimator. """ enabled: bool = True @@ -1945,6 +1727,161 @@ def collate_mixed_samples(samples: list[dict[str, Any]]) -> dict[str, Any]: return batch +class HybridJointDataLoader: + """Select one child loader for each rank-owned payload. + + ``context_parallel_size`` defines the repeating CP owner slots. Ranks + with the same slot in different replica groups select the same child for a + given fetch, while slots within one group may select different children. + A size of one synchronizes child selection across every rank, so this + loader is useful with or without context parallelism. + + ``enable_async_batch_building=True`` selects and builds the next payload + on one background thread. Keep child batch building synchronous to retain + only one future packed batch across the entire hybrid loader. Child worker + initialization and prewarming still happen on the calling thread. + """ + + def __init__( + self, + dataloaders: Mapping[str, Mapping[str, Any]], + context_parallel_size: int = 1, + seed: int = 42, + enable_async_batch_building: bool = False, + async_batch_building_timeout_s: float = 1200.0, + ) -> None: + if isinstance(context_parallel_size, bool) or not isinstance(context_parallel_size, int): + raise TypeError("context_parallel_size must be an integer.") + if context_parallel_size < 1: + raise ValueError("context_parallel_size must be positive.") + if async_batch_building_timeout_s <= 0: + raise ValueError(f"async_batch_building_timeout_s must be positive, got {async_batch_building_timeout_s}.") + + self.dataset_name_list: list[str] = [] + self.dataloader_list: list[Any] = [] + ratios: list[float] = [] + for name, entry in dataloaders.items(): + if set(entry) != {"dataloader", "ratio"}: + raise ValueError(f"Invalid child config for {name!r}: {entry}") + ratio = float(entry["ratio"]) + if not math.isfinite(ratio) or ratio < 0: + raise ValueError(f"Child ratio for {name!r} must be finite and nonnegative, got {ratio}.") + if ratio == 0: + continue + self.dataset_name_list.append(str(name)) + self.dataloader_list.append(instantiate(entry["dataloader"])) + ratios.append(ratio) + if not self.dataloader_list: + raise ValueError("At least one child loader must have a positive ratio.") + + total_ratio = sum(ratios) + self.data_probs: np.ndarray = np.asarray([ratio / total_ratio for ratio in ratios], dtype=np.float64) + self.context_parallel_size: int = context_parallel_size + self.seed: int = int(seed) + global_rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + # ParallelDims puts overlay axes last, so modulo maps a global rank to + # its rank-local owner slot for CP and to slot zero when the size is one. + self.context_parallel_rank: int = global_rank % context_parallel_size + self.global_id: int = 0 + self._child_iterators: list[Iterator[dict[str, Any]]] = [] + self.enable_async_batch_building: bool = enable_async_batch_building + self.async_batch_building_timeout_s: float = float(async_batch_building_timeout_s) + self._async_batch_builder: _AsyncBatchBuilder | None = None + + lines = [ + "HybridJointDataLoader: " + f"CP slot {self.context_parallel_rank}/{self.context_parallel_size}, " + f"{len(self.dataset_name_list)} children" + ] + for name, probability in zip(self.dataset_name_list, self.data_probs, strict=True): + lines.append(f" {name}: {probability:.1%}") + log.info("\n".join(lines), rank0_only=False) + + def _selected_child_index(self, fetch_index: int) -> int: + """Return this rank slot's child for one raw fetch.""" + owner_slot = fetch_index * self.context_parallel_size + self.context_parallel_rank + rng = np.random.RandomState(self.seed + owner_slot) + return int(rng.choice(len(self.dataloader_list), p=self.data_probs)) + + def _initialize_child_iterators_once(self) -> None: + if self._child_iterators: + return + # Initialize and prewarm every child before the first training step. + # Deferring an unused child until its first draw could stall collectives. + for dataloader in self.dataloader_list: + if self.enable_async_batch_building and isinstance(dataloader, JointDataLoader): + dataloader._validate_async_child_startup() + initialize = getattr(dataloader, "_initialize_child_iterators_once", None) + if callable(initialize): + initialize() + self._child_iterators = [iter(dataloader) for dataloader in self.dataloader_list] + + def __iter__(self) -> Iterator[dict[str, Any]]: + self._initialize_child_iterators_once() + if not self.enable_async_batch_building: + return self._iter_synchronous() + return self._iter_asynchronous() + + def _iter_synchronous(self) -> Iterator[dict[str, Any]]: + """Select and fetch each payload on the caller or the sole producer thread.""" + while True: + child_index = self._selected_child_index(self.global_id) + try: + batch = next(self._child_iterators[child_index]) + except StopIteration: + return + if not isinstance(batch, dict): + child_name = self.dataset_name_list[child_index] + raise TypeError(f"Child loader {child_name!r} yielded {type(batch).__name__}; expected dict.") + self.global_id += 1 + yield batch + + def _iter_asynchronous(self) -> Iterator[dict[str, Any]]: + """Build one future payload from the next selected child.""" + if self._async_batch_builder is None: + self._async_batch_builder = _AsyncBatchBuilder( + self._iter_synchronous(), timeout_s=self.async_batch_building_timeout_s + ) + self._async_batch_builder.request_batch(self.global_id) + async_batch_builder = self._async_batch_builder + + while True: + try: + batch = async_batch_builder.get_batch() + except StopIteration: + return + async_batch_builder.request_batch(self.global_id) + yield batch + + def set_start_iteration(self, iteration: int) -> None: + """Resume the selector and each child at its rank-local fetch count.""" + if isinstance(iteration, bool) or not isinstance(iteration, int): + raise TypeError("iteration must be an integer.") + if iteration < 0: + raise ValueError("iteration must be nonnegative.") + if self._child_iterators: + raise RuntimeError("set_start_iteration must be called before child iteration starts.") + + child_iterations = [0] * len(self.dataloader_list) + for fetch_index in range(iteration): + child_iterations[self._selected_child_index(fetch_index)] += 1 + for dataloader, child_iteration in zip(self.dataloader_list, child_iterations, strict=True): + set_start_iteration = getattr(dataloader, "set_start_iteration", None) + if callable(set_start_iteration): + set_start_iteration(child_iteration) + self.global_id = iteration + + def close(self) -> None: + """Stop background packing and forward cleanup to child loaders.""" + if self._async_batch_builder is not None: + # Retain the closed builder so another iterator cannot restart it. + self._async_batch_builder.close() + for dataloader in self.dataloader_list: + close = getattr(dataloader, "close", None) + if callable(close): + close() + + class RandomMixedJointDataLoader(JointDataLoader): """Mix sources within each packed batch without altering the legacy loaders. @@ -1953,12 +1890,25 @@ class RandomMixedJointDataLoader(JointDataLoader): sample token-mix control is explicitly enabled. Random draws allow, but do not require, several sources in a batch. - Deferred candidates are kept in draw order and offered first in the next - batch. In particular its oldest candidate starts that batch, even if its cost - exceeds the time target, so lookahead cannot starve expensive samples. A token - limit still rejects individually oversized samples using the base contract. + Lookahead counts rejections separately for each source. Once a source reaches + its limit, candidate draws exclude it for the rest of that batch and continue + among the remaining sources with renormalized probabilities. All sources get + a new lookahead allowance in the next batch. This prioritizes filling batches: + source exclusions and deferred candidates can shift the realized token mix + away from the configured shares, even with token-mix control enabled. + + Every candidate draw samples a source with the current mixing probabilities, + then retries that source's oldest deferred candidate before fetching fresh + data. Deferred candidates retain per-source order, not global priority. The + first candidate in a batch may exceed the time target; a token limit still + rejects individually oversized samples using the base contract. Finite sources are retired on exhaustion; the remaining sources continue with renormalized probabilities until all sources and pending samples finish. + + ``enable_async_batch_building=True`` runs this same packing iterator on one + background thread, preparing one future CPU batch while training consumes + the current batch. Source selection, deferred samples and token-mix state + stay on that single producer thread. """ def __init__( @@ -1982,8 +1932,13 @@ def __init__( seed: int = 42, prewarm: bool = True, prewarm_concurrency: int = 1, + prewarm_rank_stagger_s: float = 0.0, lazy_initialize_child_iterators: bool = False, forkserver_preload_modules: list[str] | None = None, + enable_async_batch_building: bool = False, + async_batch_building_timeout_s: float = 1200.0, + distributed_row_index_cache: bool = True, + lidar_patch_spatial_hw: int | tuple[int, int] | None = None, ) -> None: if default_lookahead_limit < 1 or any(value < 1 for value in (lookahead_limits or {}).values()): raise ValueError("Lookahead limits must be positive.") @@ -2003,6 +1958,8 @@ def __init__( audio_sample_rate=audio_sample_rate, prewarm=prewarm, prewarm_concurrency=prewarm_concurrency, + prewarm_rank_stagger_s=prewarm_rank_stagger_s, + distributed_row_index_cache=distributed_row_index_cache, default_lookahead_limit=default_lookahead_limit, lookahead_limits=lookahead_limits, uniae_chunk_frames=uniae_chunk_frames, @@ -2010,11 +1967,14 @@ def __init__( lazy_initialize_child_iterators=lazy_initialize_child_iterators, iteration_time_budget=iteration_time_budget, forkserver_preload_modules=forkserver_preload_modules, + enable_async_batch_building=enable_async_batch_building, + async_batch_building_timeout_s=async_batch_building_timeout_s, + lidar_patch_spatial_hw=lidar_patch_spatial_hw, ) rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 self.rng: np.random.Generator = np.random.default_rng(np.random.SeedSequence([seed, rank])) config = token_mix_control or SampleTokenMixConfig(enabled=self.iteration_time_budget is not None) - # Legacy batch-named fields count observations; here each is one fresh sample. + # Legacy batch-named fields count observations; here each is one accepted sample. # Sample lengths can differ far more than packed batches, so keep this draw unclipped. self.token_mix: TokenMixController = TokenMixController( self.dataset_name_list, @@ -2024,15 +1984,38 @@ def __init__( half_life_batches=config.half_life_samples, warmup_batches_per_source=config.warmup_samples_per_source, max_weight_ratio=None, + feedback=True, ), ) self.accepted_tokens: np.ndarray = np.zeros(len(self.dataset_name_list), dtype=np.int64) self._active: np.ndarray = np.ones(len(self.dataset_name_list), dtype=bool) - self._pending: deque[_Candidate] = deque() + self._pending: list[deque[_Candidate]] = [deque() for _ in self.dataloader_list] + + def _fresh_candidate( + self, metrics: _PackingMetrics, worker_timings: list[dict[str, Any]], eligible_sources: np.ndarray + ) -> _Candidate | None: + """Draw a source and retry its pending candidates before fetching fresh data.""" + while True: + available_sources = self._active.copy() + for source, pending in enumerate(self._pending): + if pending: + available_sources[source] = True + available_sources &= eligible_sources + if not np.any(available_sources): + return None + + source = int(self.rng.choice(len(self.dataloader_list), p=self.token_mix.probabilities(available_sources))) + + pending = self._pending[source] + if pending: + # Deferred candidates are bounded per source: + # pending_count[source] <= lookahead_limits[source]. + # Therefore, total pending candidates <= sum(self.lookahead_limits). + # Retry the selected source's pending candidates before fresh draws + # so unvisited candidates cannot accumulate beyond its allowance. + metrics.from_buffer += 1 + return pending.popleft() - def _fresh_candidate(self, metrics: _PackingMetrics, worker_timings: list[dict[str, Any]]) -> _Candidate | None: - while np.any(self._active): - source = int(self.rng.choice(len(self.dataloader_list), p=self.token_mix.probabilities(self._active))) from_buffer = bool(self.buffers[source]) try: sample = self._get_next_sample(source) @@ -2044,16 +2027,26 @@ def _fresh_candidate(self, metrics: _PackingMetrics, worker_timings: list[dict[s if not from_buffer: worker_timings.append( { - "dataset_name": self.dataset_name_list[source], + "dataset_name": sample.get("task_family", self.dataset_name_list[source]), **{key: sample[key] for key in _BATCH_TIMING_KEYS if key in sample}, } ) if _sample_should_drop(sample): - self._log_drop_sample(sample, self.dataset_name_list[source]) + self._log_drop_sample(sample, sample.get("task_family", self.dataset_name_list[source])) metrics.dropped_count += 1 _extend_action_sampler_draw_counts(metrics.dropped_action_sampler_draw_counts, sample) continue + tokens, seconds = self._compute_sample_cost(sample) + if tokens <= 0: + raise ValueError("A packable sample must contain at least one token.") + + # Check whether this sample can fit even in an empty batch. Unlike a + # temporary overflow in the packing loop, exceeding the hard token + # ceiling makes it permanently unpackable, so drop rather than defer it. + # batch_started=False ignores the time target, allowing costly singletons. + # Exclude these drops from token-mix observations, but retain their sampler + # draw counts so checkpoint resume can account for discarded draws. if not self._sample_fits( num_tokens=tokens, sample_seconds=seconds, @@ -2064,29 +2057,27 @@ def _fresh_candidate(self, metrics: _PackingMetrics, worker_timings: list[dict[s metrics.dropped_count += 1 _extend_action_sampler_draw_counts(metrics.dropped_action_sampler_draw_counts, sample) continue - if tokens <= 0: - raise ValueError("A packable sample must contain at least one token.") - self.token_mix.observe(source, tokens, record_realized=False) + return _Candidate(source, sample, tokens, seconds) - return None - def __iter__(self) -> Iterator[dict[str, Any]]: - self._initialize_child_iterators_once() - while self._pending or np.any(self._active): + def _iter_synchronous(self) -> Iterator[dict[str, Any]]: + """Pack source-weighted draws on the caller or the sole async producer thread.""" + while any(self._pending) or np.any(self._active): metrics = _PackingMetrics() samples: list[dict[str, Any]] = [] - deferred: deque[_Candidate] = deque() + deferred: list[deque[_Candidate]] = [deque() for _ in self.dataloader_list] source_stats: dict[str, dict[str, int | float]] = {} worker_timings: list[dict[str, Any]] = [] - misses = 0 + eligible_sources = np.ones(len(self.dataloader_list), dtype=bool) + stop_reason = "sample_limit" + while self.max_samples_per_batch is None or len(samples) < self.max_samples_per_batch: - if self._pending: - candidate = self._pending.popleft() - metrics.from_buffer += 1 - else: - candidate = self._fresh_candidate(metrics, worker_timings) + candidate = self._fresh_candidate(metrics, worker_timings, eligible_sources) + if candidate is None: + stop_reason = "lookahead" if np.any(~eligible_sources) else "source_exhausted" break + if not self._sample_fits( num_tokens=candidate.tokens, sample_seconds=candidate.seconds, @@ -2094,13 +2085,12 @@ def __iter__(self) -> Iterator[dict[str, Any]]: packed_sample_seconds=metrics.sample_seconds, batch_started=bool(samples), ): - deferred.append(candidate) - misses += 1 - if misses >= self.lookahead_limits[candidate.source]: - break + deferred[candidate.source].append(candidate) + if len(deferred[candidate.source]) >= self.lookahead_limits[candidate.source]: + eligible_sources[candidate.source] = False continue - name = self.dataset_name_list[candidate.source] + name = candidate.sample.get("task_family", self.dataset_name_list[candidate.source]) sample = dict(candidate.sample) sample["dataset_name"] = name samples.append(sample) @@ -2108,25 +2098,65 @@ def __iter__(self) -> Iterator[dict[str, Any]]: metrics.current_sequence_length += candidate.tokens metrics.sample_seconds += candidate.seconds self.accepted_tokens[candidate.source] += candidate.tokens + self.token_mix.observe(candidate.source, candidate.tokens) stats = source_stats.setdefault(name, {"_num_tokens": 0, "_num_samples": 0, "_sample_seconds": 0.0}) stats["_num_tokens"] += candidate.tokens stats["_num_samples"] += 1 stats["_sample_seconds"] += candidate.seconds - # Existing pending candidates preceded all fresh draws. Put candidates - # skipped this batch back ahead of those we have not visited yet. - deferred.extend(self._pending) - self._pending = deferred + # Recheck observed rejections against the completed batch: smaller + # samples admitted later may have consumed the remaining budgets. + # Mark a budget as limiting only when increasing it could admit an + # observed candidate that still fits under the other budget. + time_limited = False + token_limited = False + if ( + samples + and stop_reason != "sample_limit" + and self.max_sequence_length is not None + and self.iteration_time_budget is not None + ): + for rejected in chain.from_iterable(deferred): + tokens_fit = metrics.current_sequence_length + rejected.tokens < self.max_sequence_length + time_fits = self.iteration_time_budget.has_room_for(metrics.sample_seconds, rejected.seconds) + if tokens_fit and not time_fits: + time_limited = True + if not tokens_fit and time_fits: + token_limited = True + if time_limited and token_limited: + break + + # Within each source, pending candidates preceded fresh draws. Put + # candidates skipped this batch ahead of those we have not visited yet. + for source, deferred_candidates in enumerate(deferred): + self._pending[source].extendleft(reversed(deferred_candidates)) if not samples: return + batch = collate_mixed_samples(samples) batch["_source_packing_metrics"] = source_stats # One timing record per child batch fetch, including dropped draws. # Buffered/deferred samples do not charge their worker a second time. batch["_source_worker_timings"] = worker_timings + projected_iteration_sec = self._projected_iteration_sec(metrics.sample_seconds) + # Keep the projection unrounded for threshold comparisons in the + # monitor. This metadata describes the whole packed batch, not any + # one source's share of it, and does not change admission decisions. + batch["_packing_budget"] = { + "token_limit": self.max_sequence_length, + "time_target_sec": ( + self.iteration_time_budget.iteration_time_target if self.iteration_time_budget is not None else None + ), + "projected_sec": projected_iteration_sec, + "stop_reason": stop_reason, + "time_limited": time_limited, + "token_limited": token_limited, + } metrics.attach_to( batch, - buffer_size=len(self._pending) + sum(len(buffer) for buffer in self.buffers), - projected_iteration_sec=self._projected_iteration_sec(metrics.sample_seconds), + buffer_size=sum(len(pending) for pending in self._pending) + + sum(len(buffer) for buffer in self.buffers), + projected_iteration_sec=projected_iteration_sec, ) + self.global_id += 1 yield batch diff --git a/cosmos_framework/data/generator/lance/augment_timing.py b/cosmos_framework/data/generator/lance/augment_timing.py index 1bf7f5b03..4e6dd1d5d 100644 --- a/cosmos_framework/data/generator/lance/augment_timing.py +++ b/cosmos_framework/data/generator/lance/augment_timing.py @@ -34,6 +34,15 @@ def _augmentor_step_name(aug: Any) -> str: return getattr(aug, "__name__", None) or type(aug).__name__ +class AugmentorSampleDropError(RuntimeError): + """Signal a bounded, recoverable sample drop from an augmentor pipeline.""" + + def __init__(self, message: str, *, drop_context: str, consecutive_drops: int) -> None: + super().__init__(message) + self.drop_context = drop_context + self.consecutive_drops = consecutive_drops + + class _AugmentedWithTimingIterator: """Apply one timed augmentation step without retaining its returned sample. @@ -89,13 +98,17 @@ def last_resume_cursor(self) -> int | None: """Return the source cursor after the most recent raw sample attempt.""" return self._last_resume_cursor - def _record_drop(self) -> None: + def _record_drop(self, *, source: str) -> None: self._consecutive_drops += 1 if self._max_consecutive_drops is not None and self._consecutive_drops >= self._max_consecutive_drops: - message = f"{self._drop_context} dropped {self._consecutive_drops} consecutive samples." + message = f"{self._drop_context} dropped {self._consecutive_drops} consecutive samples; {source}." if self._drop_error_hint is not None: message = f"{message} {self._drop_error_hint}" - raise RuntimeError(message) + raise AugmentorSampleDropError( + message, + drop_context=self._drop_context, + consecutive_drops=self._consecutive_drops, + ) def __next__(self) -> dict[str, Any] | None: while True: @@ -105,7 +118,7 @@ def __next__(self) -> dict[str, Any] | None: if sample is None: self._last_resume_cursor = None self._rejection_progress.record("source", time.monotonic() - t_start, "") - self._record_drop() + self._record_drop(source="source yielded None") if self._yield_drops: return None continue @@ -120,7 +133,7 @@ def __next__(self) -> dict[str, Any] | None: result = self._prepare(sample) if self._prepare is not None else sample if result is None: self._rejection_progress.record("prepare", time.monotonic() - t_start, sample_id) - self._record_drop() + self._record_drop(source="prepare hook returned None") if self._yield_drops: return None continue @@ -145,7 +158,7 @@ def __next__(self) -> dict[str, Any] | None: if result is None: self._rejection_progress.record(step_name, time.monotonic() - t_start, sample_id) if result is None: - self._record_drop() + self._record_drop(source=f"augmentor {_augmentor_step_name(aug)!r} returned None") if self._yield_drops: return None continue diff --git a/cosmos_framework/data/generator/local_simulation/config.py b/cosmos_framework/data/generator/local_simulation/config.py index d4a071f07..6944f510a 100644 --- a/cosmos_framework/data/generator/local_simulation/config.py +++ b/cosmos_framework/data/generator/local_simulation/config.py @@ -10,13 +10,13 @@ from omegaconf import DictConfig +from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.data.generator.local_simulation.dataloader import ( LocalSimulationDataLoader, LocalSimulationJointDataLoader, ) from cosmos_framework.data.generator.local_simulation.dataset import LocalSimulationDataset from cosmos_framework.data.generator.local_simulation.resume import LocalSimulationResumeCallback -from cosmos_framework.utils.lazy_config import LazyCall as L def make_local_dataloader_config( diff --git a/cosmos_framework/data/generator/local_simulation/dataloader.py b/cosmos_framework/data/generator/local_simulation/dataloader.py index 04203511f..4d17ccb1d 100644 --- a/cosmos_framework/data/generator/local_simulation/dataloader.py +++ b/cosmos_framework/data/generator/local_simulation/dataloader.py @@ -18,13 +18,13 @@ import torch.distributed as dist from torch.utils.data import DataLoader, DistributedSampler +from cosmos_framework.utils import log from cosmos_framework.data.generator.joint_dataloader import IterativeJointDataLoader, custom_collate_fn from cosmos_framework.data.generator.local_simulation.dataset import LocalSimulationDataset from cosmos_framework.data.generator.local_simulation.resume import ( LOCAL_SIMULATION_RESUME_KEY, consume_local_resume_state, ) -from cosmos_framework.utils import log class _ResumeDistributedSampler(DistributedSampler): diff --git a/cosmos_framework/data/generator/multiview/camera_attributes.py b/cosmos_framework/data/generator/multiview/camera_attributes.py index 6dbf896fc..e6eb7d6b6 100644 --- a/cosmos_framework/data/generator/multiview/camera_attributes.py +++ b/cosmos_framework/data/generator/multiview/camera_attributes.py @@ -104,3 +104,4 @@ "camera_right_fisheye_200fov": 9, "camera_rear_fisheye_200fov": 10, } + diff --git a/cosmos_framework/data/generator/multiview/camera_attributes_test.py b/cosmos_framework/data/generator/multiview/camera_attributes_test.py index 9bb711674..df9c5bbd7 100644 --- a/cosmos_framework/data/generator/multiview/camera_attributes_test.py +++ b/cosmos_framework/data/generator/multiview/camera_attributes_test.py @@ -38,3 +38,4 @@ def test_view_ids_match_the_rig_order() -> None: assert CAMERA_VIEW_MAPPING == {name: index for index, name in enumerate(CAMERA_ATTRIBUTES)} assert sorted(CAMERA_VIEW_MAPPING.values()) == list(range(len(CAMERA_ATTRIBUTES))) + diff --git a/cosmos_framework/data/generator/multiview/camera_geometry.py b/cosmos_framework/data/generator/multiview/camera_geometry.py index c68c74dac..585ce8f68 100644 --- a/cosmos_framework/data/generator/multiview/camera_geometry.py +++ b/cosmos_framework/data/generator/multiview/camera_geometry.py @@ -56,6 +56,8 @@ PER_FRAME_CAMERA_TO_RIG_KEY = "per_frame_camera_to_rig" SOURCE_RIG_JSON_KEY = "_rig_json" +# Rig key marking a pose-world rig, whose camera-to-rig transform is each selected frame's pose. +PER_FRAME_CAMERA_TO_RIG_KEY = "per_frame_camera_to_rig" class CameraCalibrationError(ValueError): @@ -73,6 +75,24 @@ class CameraGeometrySource: provenance: dict[str, Any] = field(default_factory=dict) +def pose_world_rig_geometry(sources: dict[str, CameraGeometrySource]) -> dict[str, Any]: + """Use the shared pose world frame as the rig, for cameras that move within a scene. + + Every camera must carry world poses in one scene-fixed frame. ``camera_to_rig`` + holds each camera's first pose only to identify rig membership; a per-frame rig + reads the selected frames' ``camera_to_world`` because these cameras have no + static extrinsic. + """ + camera_to_rig: dict[str, Matrix4] = {} + for camera, source in sources.items(): + if not source.poses: + raise ValueError(f"{camera}: a pose-world rig requires camera poses") + camera_to_rig[camera] = source.poses[min(source.poses)] + if not camera_to_rig: + raise ValueError("Sensor rig requires at least one camera") + return {"camera_to_rig": camera_to_rig, PER_FRAME_CAMERA_TO_RIG_KEY: True} + + class CameraGeometry(TypedDict): camera_keys: list[str] frame_indices: torch.Tensor # [V,T] @@ -542,7 +562,8 @@ def _finalize_or_fall_back(self, data_dict: dict[str, Any]) -> dict[str, Any] | conditioned model takes its fallback for it instead of losing the sample. The pose requirement, when set, still drops samples and logs its own drops. """ - if SOURCE_CALIBRATION_ERROR_KEY not in data_dict and data_dict.get(SOURCE_GEOMETRY_KEY) is None: + if data_dict.get(SOURCE_GEOMETRY_KEY) is None: + data_dict.pop(SOURCE_CALIBRATION_ERROR_KEY, None) data_dict.pop(SOURCE_GEOMETRY_KEY, None) data_dict.pop(SOURCE_FPS_KEY, None) data_dict.pop(SOURCE_RIG_JSON_KEY, None) diff --git a/cosmos_framework/data/generator/multiview/camera_pose_actions.py b/cosmos_framework/data/generator/multiview/camera_pose_actions.py new file mode 100644 index 000000000..461e002de --- /dev/null +++ b/cosmos_framework/data/generator/multiview/camera_pose_actions.py @@ -0,0 +1,678 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Camera-pose sidecar conversion for multiview generation training.""" + +from __future__ import annotations + +import io +import re +import tarfile +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import cache +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import torch + +from cosmos_framework.data.generator.action.utils.action_processing import ( + GLOBAL_ASINH_METRIC_GRIPPER_STATS_PATH, + ActionNormalizer, + ActionProcessingRecord, + load_action_normalization_stats, + make_camera_global_asinh_normalizer, + make_pose_action_scale_normalizer, + pad_action_to_max_dim, + resolve_action_normalization, +) +from cosmos_framework.data.generator.action.utils.pose_utils import ( + PoseConvention, + RotationConvention, + pose_abs_to_rel, +) +from cosmos_framework.data.generator.action.utils.unified_action_schema import ( + EGO_POSE, + LEFT_OPENNESS, + LEFT_WRIST_POSE, + RIGHT_OPENNESS, + RIGHT_WRIST_POSE, + UNIFIED_ACTION_DIM, +) + +# Per-sample camera-to-source-stream mapping, carried from the source through the +# augmentation chain. A sidecar's rows are stamped with the source stream ID, which a +# discovering source cannot recover from the camera name: it binds a row's own streams +# to generic view slots, so the slot a stream lands in varies from sample to sample. +SOURCE_STREAM_ID_KEY = "_camera_source_stream_id" +# Per-sample ``CameraActionSlots`` of a ``unified_wrist_slots`` source, carried through the +# augmentation chain like ``SOURCE_STREAM_ID_KEY``: which unified_v1 block each camera's pose +# rows take and the gripper openness that goes with a wrist camera. +CAMERA_ACTION_SLOTS_KEY = "_camera_action_slots" + +_EXPECTED_POSE_METADATA: dict[bytes, bytes] = { + b"schema_name": b"cosmos.multiview.camera_pose", + b"camera_axes": b"x_right_y_down_z_forward", + b"transform_direction": b"camera_to_world", + b"translation_unit": b"meter", + b"quaternion_order": b"qw_qx_qy_qz", +} +_POSE_FIELDS: tuple[str, ...] = ( + "stream_id", + "sample_index", + "position_x_m", + "position_y_m", + "position_z_m", + "orientation_qw", + "orientation_qx", + "orientation_qy", + "orientation_qz", +) +_MADS_POSE_MEMBER = re.compile(r"\.(?P\d+)\.pose\.(?Pcamera_.+)\.npy$") +CameraActionNormalization = Literal["scale", "global_asinh"] +# ``per_view_ego``: every view's raw ``xyz + rot6d`` at slots ``0:9``. ``unified_wrist_slots``: a +# wrist camera's rows at its arm's unified_v1 wrist block plus that gripper's openness, every +# other camera's at the ego block ``0:9``. +CameraActionLayout = Literal["per_view_ego", "unified_wrist_slots"] +CameraActionSlot = Literal["ego", "right_wrist", "left_wrist"] +_POSE_SLICE_BY_SLOT: dict[CameraActionSlot, slice] = { + "ego": EGO_POSE, + "right_wrist": RIGHT_WRIST_POSE, + "left_wrist": LEFT_WRIST_POSE, +} +_OPENNESS_INDEX_BY_SLOT: dict[CameraActionSlot, int] = {"right_wrist": RIGHT_OPENNESS, "left_wrist": LEFT_OPENNESS} + + +@dataclass(frozen=True) +class CameraActionSlotBinding: + """The unified_v1 block one camera's pose rows take, with the gripper openness of a wrist camera. + + ``openness`` is canonical openness (``0`` closed, ``1`` fully open) per source frame, indexed like + the pose sidecar's ``sample_index``; action row ``t`` (the move from ``frame_indices[t]`` to + ``frame_indices[t + 1]``) reads ``openness[frame_indices[t + openness_frame_offset]]``, so ``1`` + pairs a row with its target frame and ``0`` with its source frame, as each single-view Action + adapter pairs its gripper column. A wrist slot requires ``openness``; the ego slot has none. + """ + + slot: CameraActionSlot + openness: np.ndarray | None = None # [F] + openness_frame_offset: int = 1 + + def __post_init__(self) -> None: + if self.slot not in _POSE_SLICE_BY_SLOT: + raise ValueError(f"Unknown camera action slot {self.slot!r}; expected {sorted(_POSE_SLICE_BY_SLOT)}") + if (self.openness is None) != (self.slot == "ego"): + raise ValueError(f"Camera action slot {self.slot!r}: openness is set exactly for a wrist slot") + if self.openness_frame_offset not in (0, 1): + raise ValueError(f"openness_frame_offset must be 0 or 1, got {self.openness_frame_offset}") + if self.openness is not None: + openness = np.asarray(self.openness) + if openness.ndim != 1 or openness.size == 0 or not np.isfinite(openness).all(): + raise ValueError(f"openness must be a nonempty finite [F] array, got shape {openness.shape}") + + +@dataclass(frozen=True) +class CameraActionSlots: + """One sample's ``CameraActionSlotBinding`` per camera key and its gripper's fully-open width. + + ``fully_open_gripper_width_m`` is the dataset's verified physical width ``W``: the canonical + metric-gripper profile's meter bounds are divided by it, exactly as the single-view Lance + adapters do, so one openness value means the same physical opening across robots. + """ + + by_camera: Mapping[str, CameraActionSlotBinding] + fully_open_gripper_width_m: float | None + + def __post_init__(self) -> None: + width = self.fully_open_gripper_width_m + if width is not None and (not np.isfinite(width) or width <= 0.0): + raise ValueError(f"fully_open_gripper_width_m must be finite and positive, got {width}") + if width is None and any(binding.slot != "ego" for binding in self.by_camera.values()): + raise ValueError("Wrist camera action slots require fully_open_gripper_width_m") + + +@dataclass(frozen=True) +class CameraPoseActionTensors: + """Camera-major action tensors derived from sampled multiview poses.""" + + action: torch.Tensor # [V*(T-1),D_model] + action_raw: torch.Tensor # [V*(T-1),D_raw] + raw_action_dim: torch.Tensor # [] + action_valid_mask: torch.Tensor # [D_model] + action_processing_record: ActionProcessingRecord + + +def read_camera_pose_table_from_bytes(payload: bytes | bytearray | memoryview) -> pa.Table: + """Read one canonical camera-pose parquet sidecar from bytes.""" + payload_bytes = bytes(payload) + if not payload_bytes: + raise ValueError("camera pose sidecar payload is empty") + return pq.read_table(pa.BufferReader(payload_bytes)) + + +def _quaternion_wxyz_from_rotation_matrix(rotation: np.ndarray) -> tuple[float, float, float, float]: + """Convert one proper 3x3 rotation matrix to a normalized wxyz quaternion.""" + trace = float(np.trace(rotation)) + if trace > 0.0: + scale = float(np.sqrt(trace + 1.0) * 2.0) + quaternion = ( + 0.25 * scale, + (float(rotation[2, 1]) - float(rotation[1, 2])) / scale, + (float(rotation[0, 2]) - float(rotation[2, 0])) / scale, + (float(rotation[1, 0]) - float(rotation[0, 1])) / scale, + ) + else: + diagonal_index = int(np.argmax(np.diag(rotation))) + if diagonal_index == 0: + scale = float(np.sqrt(1.0 + rotation[0, 0] - rotation[1, 1] - rotation[2, 2]) * 2.0) + quaternion = ( + (float(rotation[2, 1]) - float(rotation[1, 2])) / scale, + 0.25 * scale, + (float(rotation[0, 1]) + float(rotation[1, 0])) / scale, + (float(rotation[0, 2]) + float(rotation[2, 0])) / scale, + ) + elif diagonal_index == 1: + scale = float(np.sqrt(1.0 + rotation[1, 1] - rotation[0, 0] - rotation[2, 2]) * 2.0) + quaternion = ( + (float(rotation[0, 2]) - float(rotation[2, 0])) / scale, + (float(rotation[0, 1]) + float(rotation[1, 0])) / scale, + 0.25 * scale, + (float(rotation[1, 2]) + float(rotation[2, 1])) / scale, + ) + else: + scale = float(np.sqrt(1.0 + rotation[2, 2] - rotation[0, 0] - rotation[1, 1]) * 2.0) + quaternion = ( + (float(rotation[1, 0]) - float(rotation[0, 1])) / scale, + (float(rotation[0, 2]) + float(rotation[2, 0])) / scale, + (float(rotation[1, 2]) + float(rotation[2, 1])) / scale, + 0.25 * scale, + ) + norm = float(np.linalg.norm(quaternion)) + if not np.isfinite(norm) or norm <= 0.0: + raise ValueError("camera pose rotation produced an invalid quaternion") + normalized = tuple(float(value / norm) for value in quaternion) + if normalized[0] < 0.0: + normalized = tuple(-value for value in normalized) + return normalized + + +def _camera_pose_table_from_mads_frames( + stream_id: str, + frames: dict[int, np.ndarray], + *, + source_fps: float, +) -> pa.Table: + """Convert one MADS camera's contiguous OpenCV camera-to-world matrices.""" + if not frames: + raise ValueError(f"MADS pose archive has no frames for {stream_id!r}") + expected_indices = list(range(max(frames) + 1)) + if sorted(frames) != expected_indices: + raise ValueError(f"MADS pose frames for {stream_id!r} must be contiguous from 0") + + rows: list[dict[str, int | float | str]] = [] + for sample_index in expected_indices: + matrix = frames[sample_index] # [4,4] + label = f"MADS pose {stream_id} frame {sample_index}" + if matrix.shape != (4, 4) or not np.isfinite(matrix).all(): + raise ValueError(f"{label} must be a finite 4x4 matrix") + if not np.allclose(matrix[3], np.asarray([0.0, 0.0, 0.0, 1.0]), atol=1e-6): + raise ValueError(f"{label} must be homogeneous") + rotation = matrix[:3, :3] # [3,3] + if not np.allclose(rotation.T @ rotation, np.eye(3), atol=1e-4): + raise ValueError(f"{label} rotation must be orthonormal") + if not np.isclose(np.linalg.det(rotation), 1.0, atol=1e-4): + raise ValueError(f"{label} rotation must have determinant +1") + quaternion = _quaternion_wxyz_from_rotation_matrix(rotation) + rows.append( + { + "stream_id": stream_id, + "sample_index": sample_index, + "episode_time_ns": round(sample_index * 1_000_000_000 / source_fps), + "position_x_m": float(matrix[0, 3]), + "position_y_m": float(matrix[1, 3]), + "position_z_m": float(matrix[2, 3]), + "orientation_qw": quaternion[0], + "orientation_qx": quaternion[1], + "orientation_qy": quaternion[2], + "orientation_qz": quaternion[3], + } + ) + schema = pa.schema( + [ + pa.field("stream_id", pa.string(), nullable=False), + pa.field("sample_index", pa.int64(), nullable=False), + pa.field("episode_time_ns", pa.int64(), nullable=False), + pa.field("position_x_m", pa.float64(), nullable=False), + pa.field("position_y_m", pa.float64(), nullable=False), + pa.field("position_z_m", pa.float64(), nullable=False), + pa.field("orientation_qw", pa.float64(), nullable=False), + pa.field("orientation_qx", pa.float64(), nullable=False), + pa.field("orientation_qy", pa.float64(), nullable=False), + pa.field("orientation_qz", pa.float64(), nullable=False), + ], + metadata={ + **_EXPECTED_POSE_METADATA, + b"schema_version": b"1.0.0", + b"coordinate_system": b"right_handed", + }, + ) + return pa.Table.from_pylist(rows, schema=schema) + + +def camera_pose_tables_from_mads_archive( + payload: bytes, + *, + camera_keys: Sequence[str], + stream_id_mapping: Mapping[str, str], + source_fps: float, +) -> dict[str, pa.Table]: + """Convert a MADS per-frame NPY pose tar into canonical per-camera tables.""" + if not payload: + raise ValueError("MADS pose archive payload is empty") + if not np.isfinite(source_fps) or source_fps <= 0.0: + raise ValueError(f"source_fps must be positive and finite, got {source_fps}") + + wanted = set(camera_keys) + frames_by_stream: dict[str, dict[int, np.ndarray]] = {} + with tarfile.open(fileobj=io.BytesIO(payload)) as archive: + for member in archive.getmembers(): + match = _MADS_POSE_MEMBER.search(member.name) + if match is None: + continue + stream_id = match.group("stream_id") + if stream_id not in wanted: + continue + stream = archive.extractfile(member) + if stream is None: + raise ValueError(f"MADS pose archive member {member.name!r} is not a regular file") + frame_index = int(match.group("frame")) + frames = frames_by_stream.setdefault(stream_id, {}) + if frame_index in frames: + raise ValueError(f"MADS pose archive has duplicate frame {frame_index} for {stream_id!r}") + frames[frame_index] = np.load(io.BytesIO(stream.read()), allow_pickle=False) # [4,4] + + missing = [camera_key for camera_key in camera_keys if camera_key not in frames_by_stream] + if missing: + raise ValueError(f"MADS pose archive is missing camera streams {missing}") + missing_stream_ids = [camera_key for camera_key in camera_keys if camera_key not in stream_id_mapping] + if missing_stream_ids: + raise ValueError(f"MADS stream ID mapping is missing camera keys {missing_stream_ids}") + return { + camera_key: _camera_pose_table_from_mads_frames( + stream_id_mapping[camera_key], + frames_by_stream[camera_key], + source_fps=source_fps, + ) + for camera_key in camera_keys + } + + +def write_camera_pose_table(table: pa.Table, path: Path) -> None: + """Persist one canonical camera-pose table.""" + path.parent.mkdir(parents=True, exist_ok=True) + pq.write_table(table, path) + + +def _validate_pose_table_schema(table: pa.Table) -> None: + """Validate the canonical pose sidecar fields and metadata.""" + missing_fields = [field for field in _POSE_FIELDS if field not in table.column_names] + if missing_fields: + raise ValueError(f"camera pose sidecar is missing fields {missing_fields}") + + metadata = table.schema.metadata or {} + mismatched_metadata = { + key: metadata.get(key) for key, expected in _EXPECTED_POSE_METADATA.items() if metadata.get(key) != expected + } + if mismatched_metadata: + expected = {key.decode("utf-8"): value.decode("utf-8") for key, value in _EXPECTED_POSE_METADATA.items()} + got = { + key.decode("utf-8"): (value.decode("utf-8") if value is not None else None) + for key, value in mismatched_metadata.items() + } + raise ValueError(f"camera pose sidecar metadata mismatch: got {got}, expected {expected}") + + +def _as_float(row: dict[str, Any], key: str, *, row_label: str) -> float: + """Return one finite float field from a pose row.""" + try: + value = float(row[key]) + except (KeyError, TypeError, ValueError) as error: + raise ValueError(f"{row_label} has invalid {key!r}") from error + if not np.isfinite(value): + raise ValueError(f"{row_label} has non-finite {key!r}: {value}") + return value + + +def _c2w_from_pose_row(row: dict[str, Any], *, row_label: str) -> np.ndarray: + """Reconstruct one homogeneous camera-to-world matrix from a pose row.""" + qw = _as_float(row, "orientation_qw", row_label=row_label) + qx = _as_float(row, "orientation_qx", row_label=row_label) + qy = _as_float(row, "orientation_qy", row_label=row_label) + qz = _as_float(row, "orientation_qz", row_label=row_label) + quaternion = np.asarray([qw, qx, qy, qz], dtype=np.float64) # [4] + norm = float(np.linalg.norm(quaternion)) + if not np.isfinite(norm) or abs(norm - 1.0) > 1e-4: + raise ValueError(f"{row_label} has non-unit quaternion norm {norm}") + quaternion = quaternion / norm # [4] + + w, x, y, z = quaternion + xx = x * x + yy = y * y + zz = z * z + xy = x * y + xz = x * z + yz = y * z + wx = w * x + wy = w * y + wz = w * z + + matrix = np.eye(4, dtype=np.float32) # [4,4] + matrix[:3, :3] = np.asarray( + [ + [1.0 - 2.0 * (yy + zz), 2.0 * (xy - wz), 2.0 * (xz + wy)], + [2.0 * (xy + wz), 1.0 - 2.0 * (xx + zz), 2.0 * (yz - wx)], + [2.0 * (xz - wy), 2.0 * (yz + wx), 1.0 - 2.0 * (xx + yy)], + ], + dtype=np.float32, + ) # [3,3] + matrix[:3, 3] = np.asarray( + [ + _as_float(row, "position_x_m", row_label=row_label), + _as_float(row, "position_y_m", row_label=row_label), + _as_float(row, "position_z_m", row_label=row_label), + ], + dtype=np.float32, + ) # [3] + return matrix # [4,4] + + +def camera_to_world_from_pose_table( + table: pa.Table, + frame_indices: Sequence[int], + expected_stream_id: str | None = None, +) -> np.ndarray: + """Select sampled rows and reconstruct canonical camera-to-world matrices.""" + if not frame_indices: + raise ValueError("frame_indices must not be empty") + _validate_pose_table_schema(table) + if table.num_rows == 0: + raise ValueError("camera pose sidecar contains no rows") + + rows_by_index: dict[int, dict[str, Any]] = {} + for row in table.select(_POSE_FIELDS).to_pylist(): + try: + sample_index = int(row["sample_index"]) + except (KeyError, TypeError, ValueError) as error: + raise ValueError("camera pose sidecar has invalid sample_index") from error + if sample_index in rows_by_index: + raise ValueError(f"camera pose sidecar has duplicate sample_index {sample_index}") + rows_by_index[sample_index] = row + + selected_rows: list[dict[str, Any]] = [] + missing_indices: list[int] = [] + for frame_index in frame_indices: + row = rows_by_index.get(int(frame_index)) + if row is None: + missing_indices.append(int(frame_index)) + continue + if expected_stream_id is not None and str(row["stream_id"]) != expected_stream_id: + raise ValueError( + f"camera pose sidecar stream_id {row['stream_id']!r} does not match expected " + f"{expected_stream_id!r} at sample_index {frame_index}" + ) + selected_rows.append(row) + if missing_indices: + raise ValueError(f"camera pose sidecar is missing pose rows for frame indices {missing_indices}") + + poses = [ + _c2w_from_pose_row(row, row_label=f"stream {row['stream_id']} sample_index {row['sample_index']}") + for row in selected_rows + ] + return np.stack(poses, axis=0).astype(np.float32, copy=False) # [T,4,4] + + +@cache +def _make_camera_pose_action_normalizer( + *, + raw_action_dim: int, + pose_convention: PoseConvention, + camera_action_normalization: CameraActionNormalization, + camera_translation_scale: float, + camera_rotation_scale: float, +) -> ActionNormalizer: + """Build the model-space camera action normalizer used by training and inference.""" + if camera_action_normalization == "global_asinh": + if camera_translation_scale != 1.0: + raise ValueError("global_asinh camera actions require camera_translation_scale=1.0") + if camera_rotation_scale != 1.0: + raise ValueError("global_asinh camera actions require camera_rotation_scale=1.0") + return make_camera_global_asinh_normalizer(pose_convention=pose_convention) + return make_pose_action_scale_normalizer( + raw_action_dim, + translation_scale=camera_translation_scale, + rotation_scale=camera_rotation_scale, + ) + + +@cache +def make_camera_unified_global_asinh_normalizer( + *, + pose_convention: PoseConvention, + fully_open_gripper_width_m: float | None, +) -> ActionNormalizer: + """The canonical 59-D metric-gripper ``global_asinh`` normalizer in one gripper's openness units. + + The openness slots' meter bounds ``[0, R]`` become ``[0, R / W]``, as in + ``LanceActionSampler._metric_gripper_stats``; every pose slot keeps the profile's bounds. + """ + stats = load_action_normalization_stats( + "global_asinh", + stats_path=GLOBAL_ASINH_METRIC_GRIPPER_STATS_PATH, + expected_dim=UNIFIED_ACTION_DIM, + ) + if fully_open_gripper_width_m is not None: + gripper_slots = [RIGHT_OPENNESS, LEFT_OPENNESS] + q01 = stats["q01"].clone() # [59] + q99 = stats["q99"].clone() # [59] + q01[gripper_slots] /= fully_open_gripper_width_m + q99[gripper_slots] /= fully_open_gripper_width_m + stats = {**stats, "q01": q01, "q99": q99} + return resolve_action_normalization("global_asinh", stats, pose_convention=pose_convention) + + +# eq=False: generated equality would compare the mask tensor elementwise. +@dataclass(frozen=True, eq=False) +class CameraSlotActionCodec: + """Decode camera-major unified-slot actions, restoring each row's unfilled slots to exactly zero. + + The record's ``action_valid_mask`` is the union over views, so it cannot tell a head row from a + wrist row; denormalizing an unfilled slot's zero would otherwise yield a nonzero raw value. + """ + + action_normalizer: ActionNormalizer + row_valid_mask: torch.Tensor # [V*(T-1),59] bool + + def decode_to_unified( + self, normalized_action: torch.Tensor + ) -> torch.Tensor: # normalized_action: [...,V*(T-1),59], returns [...,V*(T-1),59] + if tuple(normalized_action.shape[-2:]) != tuple(self.row_valid_mask.shape): + raise ValueError( + f"camera action rows {tuple(normalized_action.shape[-2:])} do not match the recorded slot mask " + f"{tuple(self.row_valid_mask.shape)}" + ) + raw = self.action_normalizer.denormalize_action(normalized_action) # [...,V*(T-1),59] + mask = self.row_valid_mask.to(device=raw.device) # [V*(T-1),59] + return torch.where(mask, raw, torch.zeros_like(raw)) # [...,V*(T-1),59] + + +def build_unified_camera_slot_action( + per_view_actions: Sequence[np.ndarray | None], # V x [T-1,9]; None = no measured pose + *, + frame_indices: Sequence[int], + action_slots: CameraActionSlots, + camera_keys: Sequence[str], + pose_convention: PoseConvention, + max_action_dim: int, +) -> CameraPoseActionTensors: + """Scatter each view's ``xyz + rot6d`` rows into its unified_v1 block, with a wrist's openness.""" + num_rows = len(frame_indices) - 1 + num_views = len(per_view_actions) + action_raw = torch.zeros(num_views * num_rows, UNIFIED_ACTION_DIM, dtype=torch.float32) # [V*(T-1),59] + row_valid = torch.zeros(num_views * num_rows, UNIFIED_ACTION_DIM, dtype=torch.bool) # [V*(T-1),59] + frame_index_array = np.asarray(frame_indices, dtype=np.int64) # [T] + for view, (rel_pose, camera_key) in enumerate(zip(per_view_actions, camera_keys, strict=True)): + if rel_pose is None: + # Keep the camera's row range, without inventing motion for an + # unposed view. The caller excludes these rows from Action loss. + continue + binding = action_slots.by_camera.get(camera_key) + if binding is None: + raise ValueError(f"camera {camera_key!r} has no camera action slot") + if rel_pose.shape[1] != EGO_POSE.stop - EGO_POSE.start: + raise ValueError(f"unified camera action slots take xyz + rot6d rows, got width {rel_pose.shape[1]}") + rows = slice(view * num_rows, (view + 1) * num_rows) + pose_slice = _POSE_SLICE_BY_SLOT[binding.slot] + action_raw[rows, pose_slice] = torch.from_numpy(np.ascontiguousarray(rel_pose)) # [T-1,9] + row_valid[rows, pose_slice] = True + if binding.openness is not None: + openness = np.asarray(binding.openness, dtype=np.float32) # [F] + offset = binding.openness_frame_offset + openness_frames = frame_index_array[offset : offset + num_rows] # [T-1] + if openness_frames.min() < 0 or openness_frames.max() >= openness.size: + raise ValueError( + f"camera {camera_key!r}: frame indices reach {int(openness_frames.max())}, " + f"openness has {openness.size} frames" + ) + openness_index = _OPENNESS_INDEX_BY_SLOT[binding.slot] + action_raw[rows, openness_index] = torch.from_numpy(openness[openness_frames]) # [T-1] + row_valid[rows, openness_index] = True + if not torch.isfinite(action_raw).all(): + raise ValueError("camera pose action contains non-finite values") + action_normalizer = make_camera_unified_global_asinh_normalizer( + pose_convention=pose_convention, + fully_open_gripper_width_m=action_slots.fully_open_gripper_width_m, + ) + action_model = torch.where( # [V*(T-1),59] + row_valid, action_normalizer.normalize_action(action_raw), torch.zeros_like(action_raw) + ) + valid_mask_raw = row_valid.any(dim=0) # [59] + return CameraPoseActionTensors( + action=pad_action_to_max_dim(action_model, max_action_dim), # [V*(T-1),D_model] + action_raw=action_raw, + raw_action_dim=torch.tensor(UNIFIED_ACTION_DIM, dtype=torch.long), # [] + action_valid_mask=pad_action_to_max_dim(valid_mask_raw, max_action_dim).to(dtype=torch.bool), # [D_model] + action_processing_record=ActionProcessingRecord( + raw_action_dim=UNIFIED_ACTION_DIM, + action_normalizer=None, + action_valid_mask=valid_mask_raw, + action_codec=CameraSlotActionCodec(action_normalizer, row_valid.clone()), + ), + ) + + +def build_camera_pose_action( + pose_tables: Sequence[pa.Table], + *, + frame_indices: Sequence[int], + stream_ids: Sequence[str], + rotation_format: RotationConvention = "rot6d", + pose_convention: PoseConvention = "backward_chunk_anchored_16f", + max_action_dim: int = 64, + camera_action_normalization: CameraActionNormalization = "global_asinh", + camera_translation_scale: float = 1.0, + camera_rotation_scale: float = 1.0, + action_slots: CameraActionSlots | None = None, + camera_keys: Sequence[str] | None = None, +) -> CameraPoseActionTensors: + """Build camera-major action rows from sampled absolute camera poses. + + Without ``action_slots`` every view's raw ``xyz + rot6d`` sits at slots ``0:9`` (``per_view_ego``). + With them (``unified_wrist_slots``, keyed by ``camera_keys`` in view order) each view's rows take + its camera's unified_v1 block: a wrist camera's at its arm's wrist block plus that gripper's + openness, any other camera's at ``0:9``. Rows are then normalized by the canonical 59-D + ``global_asinh`` profile and every slot a view does not fill stays zero, so ``action_raw`` is + canonical 59-D and ``action_valid_mask`` is the union of the views' slots. + """ + if len(frame_indices) < 2: + raise ValueError("at least two frame indices are required to build camera pose actions") + if len(pose_tables) != len(stream_ids): + raise ValueError(f"pose table count {len(pose_tables)} does not match stream id count {len(stream_ids)}") + if max_action_dim <= 0: + raise ValueError(f"max_action_dim must be positive, got {max_action_dim}") + if action_slots is not None: + if camera_keys is None or len(camera_keys) != len(pose_tables): + raise ValueError("unified camera action slots need one camera key per pose table") + if camera_action_normalization != "global_asinh": + raise ValueError("unified camera action slots require camera_action_normalization='global_asinh'") + if rotation_format != "rot6d": + raise ValueError("unified camera action slots require rotation_format='rot6d'") + if camera_translation_scale != 1.0 or camera_rotation_scale != 1.0: + raise ValueError("global_asinh camera actions require unit camera translation and rotation scales") + + per_view_actions: list[np.ndarray] = [] + raw_action_dim: int | None = None + for pose_table, stream_id in zip(pose_tables, stream_ids, strict=True): + c2w = camera_to_world_from_pose_table( # [T,4,4] + pose_table, + frame_indices, + expected_stream_id=stream_id, + ) + rel_pose = pose_abs_to_rel( # [T-1,D_raw] + c2w, + rotation_format=rotation_format, + pose_convention=pose_convention, + ) + if rel_pose.ndim != 2 or rel_pose.shape[0] != len(frame_indices) - 1: + raise ValueError( + f"relative camera poses for stream {stream_id!r} have shape {rel_pose.shape}, " + f"expected ({len(frame_indices) - 1}, D)" + ) + if raw_action_dim is None: + raw_action_dim = int(rel_pose.shape[1]) + elif rel_pose.shape[1] != raw_action_dim: + raise ValueError( + f"relative camera pose width for stream {stream_id!r} is {rel_pose.shape[1]}, expected {raw_action_dim}" + ) + per_view_actions.append(rel_pose.astype(np.float32, copy=False)) + + assert raw_action_dim is not None + if action_slots is not None: + assert camera_keys is not None + return build_unified_camera_slot_action( + per_view_actions, + frame_indices=frame_indices, + action_slots=action_slots, + camera_keys=camera_keys, + pose_convention=pose_convention, + max_action_dim=max_action_dim, + ) + action_raw_np = np.concatenate(per_view_actions, axis=0).astype(np.float32, copy=False) # [V*(T-1),D_raw] + action_raw = torch.from_numpy(np.ascontiguousarray(action_raw_np)).float() # [V*(T-1),D_raw] + if not torch.isfinite(action_raw).all(): + raise ValueError("camera pose action contains non-finite values") + valid_mask_raw = torch.ones(raw_action_dim, dtype=torch.bool) # [D_raw] + action_normalizer = _make_camera_pose_action_normalizer( + raw_action_dim=raw_action_dim, + pose_convention=pose_convention, + camera_action_normalization=camera_action_normalization, + camera_translation_scale=camera_translation_scale, + camera_rotation_scale=camera_rotation_scale, + ) + action_model = action_normalizer.normalize_action(action_raw) # [V*(T-1),D_raw] + action = pad_action_to_max_dim(action_model, max_action_dim) # [V*(T-1),D_model] + action_valid_mask = pad_action_to_max_dim(valid_mask_raw, max_action_dim).to(dtype=torch.bool) # [D_model] + action_processing_record = ActionProcessingRecord( + raw_action_dim=raw_action_dim, + action_normalizer=action_normalizer, + action_valid_mask=valid_mask_raw, + ) + return CameraPoseActionTensors( + action=action, + action_raw=action_raw, + raw_action_dim=torch.tensor(raw_action_dim, dtype=torch.long), # [] + action_valid_mask=action_valid_mask, + action_processing_record=action_processing_record, + ) diff --git a/cosmos_framework/data/generator/multiview/caption_format.py b/cosmos_framework/data/generator/multiview/caption_format.py index 3d13d3100..34d125c21 100644 --- a/cosmos_framework/data/generator/multiview/caption_format.py +++ b/cosmos_framework/data/generator/multiview/caption_format.py @@ -35,6 +35,14 @@ _TIMESTAMP: re.Pattern[str] = re.compile(r"\[\s*\d+(?:\.\d+)?\s*s?\s*[-–—]\s*\d+(?:\.\d+)?\s*s?\s*\]", re.IGNORECASE) +def format_short_task_caption(caption: str) -> str: + """Express a short task instruction as a goal before adding view metadata.""" + caption = caption.strip() + if not caption or caption.startswith("The goal is "): + return caption + return f"The goal is {caption[0].lower()}{caption[1:]}" + + def first_caption_paragraph(caption: str) -> str: """Keep the opening timestamped paragraph, or the first blank-line paragraph.""" caption = caption.strip() @@ -77,14 +85,23 @@ def format_separate_view_captions( camera_names: Sequence[str], add_camera_rig_prefix: bool = False, camera_attributes: Mapping[str, Mapping[str, str | int]] | None = None, + append_camera_role_to_caption: bool = False, + camera_prefixes: Mapping[str, str] | None = None, ) -> list[Caption]: - """Preserve per-view captions, optionally adding the sampled rig in natural language.""" + """Preserve per-view captions, optionally adding physical camera roles and the sampled rig.""" if not captions: raise ValueError("Separate-view captions must not be empty") if len(captions) != len(camera_names): raise ValueError( f"Separate captions must match sampled cameras: captions={len(captions)}, cameras={len(camera_names)}" ) + if append_camera_role_to_caption: + captions = [ + f"{camera_prefixes[camera_name]}\n\n{caption if isinstance(caption, str) else json.dumps(caption)}" + if camera_prefixes and camera_name in camera_prefixes + else caption + for camera_name, caption in zip(camera_names, captions) + ] if not add_camera_rig_prefix: return list(captions) if camera_attributes is None: diff --git a/cosmos_framework/data/generator/processors/qwen3vl_processor.py b/cosmos_framework/data/generator/processors/qwen3vl_processor.py index f4c448153..f1db77295 100644 --- a/cosmos_framework/data/generator/processors/qwen3vl_processor.py +++ b/cosmos_framework/data/generator/processors/qwen3vl_processor.py @@ -3,7 +3,7 @@ import json from pathlib import Path -from typing import Optional +from typing import Any, Optional import numpy as np import torch @@ -14,6 +14,16 @@ maybe_get_max_pixels_from_images_kwargs, maybe_parse_video_content, ) +from cosmos_framework.utils.generator.structured_conversation import ( + REASONING_PREFILL_KEY, + assistant_body_spans, + restore_agentic_messages, + restrict_assistant_loss, +) + +# Set by the reasoner video decoder on media (and copied onto chat content) whose frames it already +# smart-resized to the sample's video token budget. +VIDEO_BUDGET_RESIZED_KEY = "budget_resized" def _is_local_qwen35_snapshot(path: str) -> bool: @@ -52,9 +62,7 @@ def __init__( self.merge_size = self.processor.video_processor.merge_size self.use_smart_resize = True normalized_name = f"{name} {type(self.processor).__name__}".lower() - self.retain_mm_token_type_ids: bool = ( - "qwen3.5" in normalized_name or "qwen3_5" in normalized_name or _is_local_qwen35_snapshot(name) - ) + self.retain_mm_token_type_ids: bool = "qwen3.5" in normalized_name or "qwen3_5" in normalized_name or _is_local_qwen35_snapshot(name) def apply_chat_template( self, @@ -76,8 +84,15 @@ def apply_chat_template( assert tokenize, "tokenize must be True" assert return_tensors == "pt", "return_tensors must be pt" # Note: this tokenizer does not support "content": str, it always expect "content" entry to be a list of dicts + messages, stored_tools = restore_agentic_messages(messages) messages = convert_string_content_to_list_content(messages) - kwargs = {} + tools = kwargs.pop("tools", stored_tools) + if stored_tools is not None and tools != stored_tools: + raise ValueError("Explicit tools disagree with the stored tool menu") + # ``add_vision_id`` is a chat-template variable ("Picture N: " / "Video N: " labels). TokenizeData requests it + # only when its opt-in ``add_vision_id`` is set, so existing checkpoints keep seeing unlabeled media; it is + # forwarded only when set because HF warns about it on calls without media. + template_kwargs = {"add_vision_id": True} if kwargs.pop("add_vision_id", False) else {} max_pixels, min_pixels = maybe_get_max_pixels_from_images_kwargs(messages) if max_pixels is not None: @@ -90,7 +105,6 @@ def apply_chat_template( # File "/invalid_dir", line 321, in _decode_and_sample_videos # raise ValueError( # ValueError: Sampling frames from a list of images is not supported! Set `do_sample_frames=False`. - kwargs["videos_kwargs"] = dict(do_sample_frames=False) video_metadata = [ dict(fps=fps, total_num_frames=total_num_frames, frames_indices=frames_indices) for fps, total_num_frames, frames_indices in zip( @@ -103,6 +117,10 @@ def apply_chat_template( "do_sample_frames": False, "video_metadata": video_metadata[0] if num_video == 1 else video_metadata, } + if self._videos_are_budget_resized(messages): + # The HF resize would re-apply video_processor.size["longest_edge"] (12,288 tokens per video for the + # released Qwen3-VL / Qwen3.5 configs) and shrink frames the decoder already sized to the budget. + kwargs["videos_kwargs"]["do_resize"] = False # Both Transformers 4.x and 5.x forward modality kwargs directly to # the processor. Nesting them in ``processor_kwargs`` silently loses # frame indices, FPS and the request to keep already-sampled frames. @@ -112,6 +130,8 @@ def apply_chat_template( add_generation_prompt=add_generation_prompt, return_dict=True, return_tensors=return_tensors, + tools=tools, + **template_kwargs, **kwargs, ) @@ -132,6 +152,28 @@ def apply_chat_template( inputs["mm_token_type_ids"] = self.build_mm_token_type_ids(inputs["input_ids"]) return inputs + def _videos_are_budget_resized(self, messages: list[dict]) -> bool: + """Whether the reasoner video decoder already sized every video of the conversation to its token budget. + + Videos from any other producer (e.g. raw clips passed by the VFM model) keep the HF resize. + + Raises: + ValueError: Only some videos are decoder-sized. HF takes one resize setting per call, so it would either + re-cap the decoder-sized videos below their token budget or leave the others unresized. + """ + flags = [ + isinstance(content.get("video"), list) and bool(content.get(VIDEO_BUDGET_RESIZED_KEY, False)) + for message in messages + if isinstance(message["content"], list) + for content in message["content"] + if content.get("type", "") == "video" + ] + if not any(flags): + return False + if not all(flags): + raise ValueError("Either every video of a conversation is sized by the reasoner decoder or none is") + return True + def build_mm_token_type_ids(self, input_ids: torch.Tensor) -> torch.Tensor: """Identify text, image and video tokens in the final expanded sequence.""" mm_token_type_ids = torch.zeros_like(input_ids, dtype=torch.long) @@ -141,7 +183,39 @@ def build_mm_token_type_ids(self, input_ids: torch.Tensor) -> torch.Tensor: mm_token_type_ids[input_ids == self.video_token_id] = 2 return mm_token_type_ids - def add_assistant_tokens_mask(self, tokens): + def apply_supervision_overrides( + self, + tokens: torch.Tensor | list[int], + mask: torch.Tensor | list[bool], + messages: list[dict[str, Any]], + ) -> torch.Tensor | list[bool]: + """Preserve plain-message masks while honoring structured source supervision.""" + structured = any( + any(key in message for key in ("tools", "tool_calls", "loss", REASONING_PREFILL_KEY)) + or message.get("role") == "tool" + for message in messages + ) + if not structured: + return mask + values = tokens.detach().cpu().numpy() if isinstance(tokens, torch.Tensor) else np.asarray(tokens) + # Structured messages may merge source-leading whitespace into the header + # token. Repair that boundary without replacing processor-specific masks + # (in particular, Nemo Chat's generation-supplied thinking prefix). + losses = ( + mask.detach().cpu().numpy().copy() + if isinstance(mask, torch.Tensor) + else np.asarray(mask, dtype=bool).copy() + ) + tokenizer = self.processor.tokenizer + frame_starts = np.flatnonzero(values == tokenizer.convert_tokens_to_ids("<|im_start|>")) + header_size = 1 + len(tokenizer.encode("assistant\n", add_special_tokens=False)) + for start, end in assistant_body_spans(values, tokenizer): + frame_start = int(frame_starts[np.searchsorted(frame_starts, start) - 1]) + losses[start : min(frame_start + header_size, end)] = True + restricted = restrict_assistant_loss(values, losses, messages, self.processor.tokenizer) + return torch.from_numpy(restricted) if isinstance(tokens, torch.Tensor) else restricted.tolist() + + def add_assistant_tokens_mask(self, tokens: list[int] | torch.Tensor) -> list[bool] | torch.Tensor: """ Add a mask to the assistant tokens. This is used to mask out tokens that are not generated by the assistant (e.g., system prompts, user prompts, chat templates), such that in the loss computation, only the tokens generated by the assistant are used. @@ -161,18 +235,21 @@ def add_assistant_tokens_mask(self, tokens): np_tokens = tokens.cpu().numpy() if isinstance(tokens, torch.Tensor) else np.array(tokens) assert np_tokens.ndim == 1 - # Constants defining bos, eos and fixed offsets. + # Constants defining the ChatML boundaries. BOS_TOKEN = "<|im_start|>" EOS_TOKEN = "<|im_end|>" ROLE = "assistant" - # Offsets: skip the bos + "assistant\n" (always 3 tokens) and include the eos (+1) for supervision - START_OFFSET = 3 END_OFFSET = 1 - # Retrieve token IDs for the markers and the role. + # Match tokenizer-rendered headers instead of assuming that the role is + # one token. Some supported tokenizers render ``assistant`` as several + # BPE tokens followed by a newline token. bos_token_id = self.processor.tokenizer.convert_tokens_to_ids(BOS_TOKEN) eos_token_id = self.processor.tokenizer.convert_tokens_to_ids(EOS_TOKEN) role_id = self.processor.tokenizer.convert_tokens_to_ids(ROLE) + full_role_header_ids = self.processor.tokenizer.encode(f"{BOS_TOKEN}{ROLE}\n", add_special_tokens=False) + role_header_ids = self.processor.tokenizer.encode(f"{ROLE}\n", add_special_tokens=False) + role_ids = self.processor.tokenizer.encode(ROLE, add_special_tokens=False) # Locate all positions where the start and end markers appear. start_indices = np.where(np_tokens == bos_token_id)[0] @@ -181,12 +258,29 @@ def add_assistant_tokens_mask(self, tokens): # Initialize the mask with False values. masks = np.zeros_like(np_tokens, dtype=bool) assert len(start_indices) == len(end_indices) - # For each pair of bos/eos, check if the role is 'assistant' - # and apply the mask accordingly. + # For each pair of bos/eos, check if the role is assistant and apply + # the mask to the response plus the closing boundary. for start, end in zip(start_indices, end_indices): - if np_tokens[start + 1] == role_id: - # Mask tokens from after the assistant header (start+3) to include the end marker (end+1) - masks[start + START_OFFSET : end + END_OFFSET] = True + content_start = None + full_rendered_header = np.asarray(full_role_header_ids, dtype=np_tokens.dtype) + header_end = start + len(full_rendered_header) + if header_end <= end and np.array_equal(np_tokens[start:header_end], full_rendered_header): + content_start = header_end + else: + rendered_header = np.asarray(role_header_ids, dtype=np_tokens.dtype) + header_end = start + 1 + len(rendered_header) + if header_end <= end and np.array_equal(np_tokens[start + 1 : header_end], rendered_header): + content_start = header_end + if content_start is None and start + 1 < end and np_tokens[start + 1] == role_id: + content_start = start + 3 + elif content_start is None: + rendered_role = np.asarray(role_ids, dtype=np_tokens.dtype) + role_end = start + 1 + len(rendered_role) + if role_end < end and np.array_equal(np_tokens[start + 1 : role_end], rendered_role): + content_start = role_end + 1 + if content_start is None or content_start > end: + continue + masks[content_start : end + END_OFFSET] = True assert masks.shape == np_tokens.shape if isinstance(tokens, torch.Tensor): diff --git a/cosmos_framework/data/generator/processors/qwen3vl_processor_test.py b/cosmos_framework/data/generator/processors/qwen3vl_processor_test.py new file mode 100644 index 000000000..38d2a8c14 --- /dev/null +++ b/cosmos_framework/data/generator/processors/qwen3vl_processor_test.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Keyword arguments the Qwen3-VL wrapper forwards to the HuggingFace processor.""" + +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +from PIL import Image + +from cosmos_framework.data.generator.processors import base +from cosmos_framework.data.generator.processors.qwen3vl_processor import VIDEO_BUDGET_RESIZED_KEY, Qwen3VLProcessor + + +class _RecordingHFProcessor: + """Stand-in for the HF Qwen3-VL processor (Qwen3-VL geometry) that records how it is called.""" + + image_token = "<|image_pad|>" + video_token = "<|video_pad|>" + + def __init__(self) -> None: + self.image_processor = SimpleNamespace(size={"shortest_edge": 65536, "longest_edge": 16777216}) + self.video_processor = SimpleNamespace(patch_size=16, temporal_patch_size=2, merge_size=2) + self.tokenizer = SimpleNamespace( + convert_tokens_to_ids={"<|image_pad|>": 1, "<|video_pad|>": 2}.get, eos_token_id=3, pad_token_id=0 + ) + self.calls: list[dict[str, Any]] = [] + + def apply_chat_template(self, messages: list[dict[str, Any]], **kwargs: Any) -> dict[str, torch.Tensor]: + self.calls.append(kwargs) + input_ids = torch.zeros(1, 4, dtype=torch.long) # [1,N] + return {"input_ids": input_ids, "attention_mask": torch.ones_like(input_ids)} # [1,N] each + + +@pytest.fixture +def processor(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Qwen3VLProcessor: + hf_processor = _RecordingHFProcessor() + monkeypatch.setattr(base.AutoProcessor, "from_pretrained", lambda *args, **kwargs: hf_processor) + return Qwen3VLProcessor(name=str(tmp_path)) + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize( + "budget_resized", + [ + True, # sized by the reasoner decoder: an HF resize would re-cap the token budget + False, # frames from any other producer keep the HF resize + ], +) +def test_hf_video_resize_is_skipped_only_for_decoder_sized_frames( + processor: Qwen3VLProcessor, budget_resized: bool +) -> None: + frames = [Image.new("RGB", (96, 64)) for _ in range(4)] + video = {"type": "video", "video": frames, "fps": 2.0, VIDEO_BUDGET_RESIZED_KEY: budget_resized} + + processor.apply_chat_template([{"role": "user", "content": [video, {"type": "text", "text": "Describe."}]}]) + + kwargs = processor.processor.calls[-1] + videos_kwargs = kwargs.get("processor_kwargs", kwargs)["videos_kwargs"] # transformers>=5 nests them + assert videos_kwargs.get("do_resize", True) is not budget_resized + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize( + "other_video", + [ + [Image.new("RGB", (96, 64)) for _ in range(4)], # frames from another producer + "clip.mp4", # a video HF loads itself + ], +) +def test_videos_mixing_decoder_sized_and_other_videos_are_rejected( + processor: Qwen3VLProcessor, other_video: list[Image.Image] | str +) -> None: + """HF takes one resize setting per call: it would re-cap the decoder-sized video or leave the other unresized.""" + frames = [Image.new("RGB", (96, 64)) for _ in range(4)] + sized = {"type": "video", "video": frames, "fps": 2.0, VIDEO_BUDGET_RESIZED_KEY: True} + other = {"type": "video", "video": other_video, "fps": 2.0} + + with pytest.raises(ValueError, match="sized by the reasoner decoder"): + processor.apply_chat_template([{"role": "user", "content": [sized, other, {"type": "text", "text": "?"}]}]) + + assert processor.processor.calls == [] + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("add_vision_id", [False, True]) +def test_add_vision_id_reaches_the_chat_template_only_when_set( + processor: Qwen3VLProcessor, add_vision_id: bool +) -> None: + messages = [{"role": "user", "content": [{"type": "text", "text": "Compare."}]}] + + processor.apply_chat_template(messages, add_vision_id=add_vision_id) + + kwargs = processor.processor.calls[-1] + assert kwargs.get("add_vision_id", False) is add_vision_id + assert "add_vision_id" not in kwargs.get("processor_kwargs", {}) # a template variable, not a processor kwarg diff --git a/cosmos_framework/data/generator/reasoner/collate_fn.py b/cosmos_framework/data/generator/reasoner/collate_fn.py index e2f060ba9..04da297be 100644 --- a/cosmos_framework/data/generator/reasoner/collate_fn.py +++ b/cosmos_framework/data/generator/reasoner/collate_fn.py @@ -312,11 +312,10 @@ def _collate_true_packing(batch: list[dict[str, Any]]) -> dict[str, Any]: over_budget: int | None = None for item in batch: item.pop("true_packing", None) - # attention_mask is unused on the true-packing forward: varlen attention is driven by - # cu_seqlens, packed position_ids are built from seq_lens, and the pad tail is kept inert by - # labels=ignore_index + token_mask=False. vlm_model pops it before the forward, so emitting a - # per-step [1, S_pad] mask would only add a wasted H2D copy (true packing). Drop it here so - # the variable-length per-sample masks also never reach the generic key handling below. + # Varlen attention is driven by cu_seqlens and packed positions by seq_lens. + # Drop the variable-length per-sample masks before generic key handling. + # Qwen's CPU preparation rebuilds one real-token mask for CP slicing and + # MoE routing; labels=ignore_index + token_mask=False mask the pad-tail loss. item.pop("attention_mask", None) v = item.pop("predicted_runtime_ms", None) sc = item.pop("singleton_cause", None) diff --git a/cosmos_framework/data/generator/reasoner/video_decoder_qwen.py b/cosmos_framework/data/generator/reasoner/video_decoder_qwen.py index 4ab7391b7..790126386 100644 --- a/cosmos_framework/data/generator/reasoner/video_decoder_qwen.py +++ b/cosmos_framework/data/generator/reasoner/video_decoder_qwen.py @@ -25,7 +25,7 @@ from torchvision.transforms import InterpolationMode from cosmos_framework.utils import log -from cosmos_framework.data.generator.processors.qwen3vl_processor import Qwen3VLProcessor +from cosmos_framework.data.generator.processors.qwen3vl_processor import VIDEO_BUDGET_RESIZED_KEY, Qwen3VLProcessor from cosmos_framework.utils.generator.source_video_timing import ( SOURCE_VIDEO_TIMING_KEY, build_source_video_timing, @@ -42,6 +42,10 @@ VideoTemporalMode = Literal["native", "framewise"] +class VideoDecodeTimeoutError(TimeoutError): + """A resource-dependent decode failure that must not become a fixed data drop.""" + + def get_effective_temporal_patch_size(temporal_patch_size: int, video_temporal_mode: VideoTemporalMode) -> int: """Return the temporal token-budget factor without changing the ViT tubelet kernel.""" if video_temporal_mode == "native": @@ -75,14 +79,17 @@ def pixels_to_token(pixels: int, patch_size: int = 14, temporal_patch_size: int def video_decoder_qwen( num_threads: int = 0, min_fps_thres: int = 4, - max_fps_thres: int = 60, + max_fps_thres: float = 60.0, target_fps: float = 2.0, min_video_token_length: int = 16, max_video_token_length: int = 8192, + max_video_frames: int | None = None, random_augmentation: bool = False, frame_count_random_range: Optional[list[int]] = None, video_timestamp_mode: str = "qwen_index", video_temporal_mode: VideoTemporalMode = "native", + decoding_timeout: int = 60, + raise_on_decode_timeout: bool = False, **kwargs, ) -> Callable: """ @@ -118,10 +125,13 @@ def video_decoder_qwen( target_fps=target_fps, min_video_token_length=min_video_token_length, max_video_token_length=max_video_token_length, + max_video_frames=max_video_frames, random_augmentation=random_augmentation, frame_count_random_range=frame_count_random_range, video_timestamp_mode=video_timestamp_mode, video_temporal_mode=video_temporal_mode, + decoding_timeout=decoding_timeout, + raise_on_decode_timeout=raise_on_decode_timeout, ) return video_decoder_configured @@ -132,10 +142,11 @@ def _video_decoder_qwen_func( data: bytes, processor: Qwen3VLProcessor, min_fps_thres: int = 4, - max_fps_thres: int = 60, + max_fps_thres: float = 60.0, target_fps: float = 2.0, min_video_token_length: int = 16, max_video_token_length: int = 8192, + max_video_frames: int | None = None, num_threads: int = 0, random_augmentation: bool = False, fps_random_range: list[float] = [0.5, 1.5], @@ -146,6 +157,8 @@ def _video_decoder_qwen_func( end_frame: Optional[int] = None, decoding_timeout: int = 60, video_timestamp_mode: str = "qwen_index", + raise_on_decode_timeout: bool = False, + decode_chunk_size: int = 32, **kwargs, ) -> dict | None: """Actual video decoder function. @@ -165,9 +178,12 @@ def _video_decoder_qwen_func( frame_count_random_range (list[int], optional): Random frame count range. If provided, take priority over fps_random_range. start_frame (Optional[int], optional): Start frame. Defaults to None. If both start_frame and end_frame are provided, the video will be decoded from start_frame to end_frame. end_frame (Optional[int], optional): End frame. Defaults to None. If both start_frame and end_frame are provided, the video will be decoded from start_frame to end_frame. - decoding_timeout (int, optional): Timeout in seconds. Defaults to 60. + decoding_timeout (int, optional): Timeout in seconds for decoding one chunk of frames. Defaults to 60. + raise_on_decode_timeout: Raise instead of dropping a sample when decoding times out. video_timestamp_mode: Use Qwen source index/FPS by default for supported processors, legacy synthetic FPS, or actual selected FrameBatch presentation timestamps. + decode_chunk_size (int, optional): Frames decoded (and resized) at a time, bounding native-resolution + frames held in memory. Defaults to 32. Raises: ValueError: Video fps lower than 1, skipping ValueError: Video fps lower than min_fps_thres, skipping @@ -252,7 +268,20 @@ def _video_decoder_qwen_func( max_pixels: int = token_to_pixels(max_video_token_length, patch_size, temporal_patch_size, merge_size) total_min_pixels = min_pixels total_max_pixels = max_pixels - max_frames: int = max_pixels // (min_height_width) ** 2 // temporal_patch_size + # A budget whose frame cap is below one temporal patch (under 128 tokens for Qwen3-VL in native mode, e.g. a + # short clip sharing a budget with long videos) still decodes one temporal patch; the per-frame budget below + # sets its resolution, under min_height_width when the budget cannot afford more. + max_frames: int = max(max_pixels // (min_height_width) ** 2 // temporal_patch_size, temporal_patch_size) + if max_video_frames is not None: + if max_video_frames <= 0: + raise ValueError(f"max_video_frames must be positive, got {max_video_frames}") + aligned_max_video_frames = max_video_frames // temporal_patch_size * temporal_patch_size + if aligned_max_video_frames < temporal_patch_size: + raise ValueError( + f"max_video_frames must allow at least {temporal_patch_size} temporally aligned frames, " + f"got {max_video_frames}" + ) + max_frames = min(max_frames, aligned_max_video_frames) nframes = smart_nframes_with_factor( {"fps": target_fps, "max_frames": max_frames}, @@ -265,64 +294,89 @@ def _video_decoder_qwen_func( else: idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist() # [nframes] - def _decode_video() -> torch.Tensor | FrameBatch: # pixels: [T,C,H,W]; times: [T] - frame_batch = video_reader.get_frames_at(indices=idx) # pixels: [T,C,H,W]; times: [T] - return frame_batch if video_timestamp_mode == "source_pts" else frame_batch.data # [T,C,H,W] uint8 - - # Use ThreadPoolExecutor to run video decoding with a timeout. - # If the thread is stuck, abandon it immediately. + # Each frame gets an even share of the budget. min_pixels floors the frame size only where that share affords + # it: below it (the single temporal patch a too-small budget still decodes, see max_frames) the floor is one + # merged patch, so frames are only ever shrunk to fit and a small allocation never spends more tokens than it + # was given. + max_pixels_per_frame = max_pixels // nframes + min_pixels_per_frame = min_pixels if max_pixels_per_frame >= min_pixels else (patch_size * merge_size) ** 2 + + # Decode (and resize) chunk by chunk into one preallocated output. With a resize, only one native chunk is alive + # at a time: decoding every sampled frame at native resolution before resizing peaks at ~37 MiB per 1080p frame, + # tens of GiB for the frame counts of long-context budgets. Without one, the output itself is native; filling it + # in place still avoids holding the decoded chunks and their concatenation at once. + # Each chunk is decoded in a worker thread with a timeout; a stuck thread is abandoned immediately. executor = ThreadPoolExecutor(max_workers=1) - future = executor.submit(_decode_video) + video_frames: torch.Tensor | None = None # [T,C,H',W'] uint8, allocated once the first chunk gives the size + pts_chunks: list[torch.Tensor] = [] + duration_chunks: list[torch.Tensor] = [] try: - decoded = future.result(timeout=decoding_timeout) # pixels: [T,C,H,W]; optional times: [T] - executor.shutdown(wait=False) + for chunk_start in range(0, nframes, decode_chunk_size): + chunk_indices = idx[chunk_start : chunk_start + decode_chunk_size] + frame_batch = executor.submit(video_reader.get_frames_at, indices=chunk_indices).result( + timeout=decoding_timeout + ) # pixels: [t,C,H,W]; times: [t] + if video_timestamp_mode == "source_pts": + if not isinstance(frame_batch, FrameBatch): + raise ValueError("source_pts requires TorchCodec FrameBatch presentation metadata") + pts_chunks.append(frame_batch.pts_seconds) # [t] + duration_chunks.append(frame_batch.duration_seconds) # [t] + frames = frame_batch.data # [t,C,H,W] uint8 + del frame_batch # the native chunk must not outlive its resize while the next chunk decodes + if frames.shape[0] != len(chunk_indices): + raise ValueError(f"[{key}] decoder returned {frames.shape[0]} frames for {len(chunk_indices)} indices") + if video_frames is None: + native_height, native_width = int(frames.shape[2]), int(frames.shape[3]) + resized_height, resized_width = native_height, native_width + if processor.use_smart_resize: + resized_height, resized_width = smart_resize( + native_height, + native_width, + factor=patch_size * merge_size, + min_pixels=min_pixels_per_frame, + max_pixels=max_pixels_per_frame, + ) + log.debug( + f"resized_height: {resized_height}, resized_width: {resized_width} | original height: {native_height}, original width: {native_width}" + ) + video_frames = torch.empty( + (nframes, frames.shape[1], resized_height, resized_width), dtype=frames.dtype + ) # [T,C,H',W'] + if processor.use_smart_resize: + frames = transforms.functional.resize( + frames, + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ) # [t,C,H',W'] uint8 + video_frames[chunk_start : chunk_start + len(chunk_indices)] = frames + del frames # copied; the chunk must not stay alive while the next one decodes except TimeoutError as e: log.warning(f"[{key}] Video decoding timed out after {decoding_timeout} seconds") - executor.shutdown(wait=False) + if raise_on_decode_timeout: + raise VideoDecodeTimeoutError( + f"Video decoding exceeded {decoding_timeout} seconds for {key!r}; " + "retry with sufficient CPU resources or a larger timeout instead of dropping this sample." + ) from e return None + finally: + executor.shutdown(wait=False) + if video_frames is None: + raise ValueError(f"[{key}] no frames were sampled") + if processor.use_smart_resize: + video_frames = video_frames.float() # [T,C,H',W'] timing = None if video_timestamp_mode == "source_pts": - if not isinstance(decoded, FrameBatch): - raise ValueError("source_pts requires TorchCodec FrameBatch presentation metadata") timing = build_source_video_timing( idx, - decoded.pts_seconds, # [T] - decoded.duration_seconds, # [T] + torch.cat(pts_chunks), # [T] + torch.cat(duration_chunks), # [T] hashlib.sha256(data).hexdigest(), ) - video_frames = decoded.data # [T,C,H,W] - else: - assert isinstance(decoded, torch.Tensor) - video_frames = decoded # [T,C,H,W] sample_fps = nframes / max(total_frames, 1e-6) * video_fps - - # recompute max_pixels based on number of sampled frames - nframes, _, height, width = video_frames.shape - native_height, native_width = int(height), int(width) - # Floor the per-frame budget at min_pixels so smart_resize never gets an inverted range - # (min_pixels > max_pixels) for many-frame videos. Matches the projects/cosmos3/vlm decoder - # the cosmos-rl i4 data bridge used. - max_pixels = max(max_pixels // nframes, min_pixels) - if processor.use_smart_resize: - resized_height, resized_width = smart_resize( - height, - width, - factor=patch_size * merge_size, - min_pixels=min_pixels, - max_pixels=max_pixels, - ) - log.debug( - f"resized_height: {resized_height}, resized_width: {resized_width} | original height: {height}, original width: {width}" - ) - video_frames = transforms.functional.resize( - video_frames, - [resized_height, resized_width], - interpolation=InterpolationMode.BICUBIC, - antialias=True, - ).float() # [T,C,H,W] - video_frames = video_frames.permute(1, 0, 2, 3) # [C,T,H,W] + video_frames = video_frames.permute(1, 0, 2, 3) # [C,T,H',W'] result = dict( videos=video_frames, @@ -338,6 +392,7 @@ def _decode_video() -> torch.Tensor | FrameBatch: # pixels: [T,C,H,W]; times: [ budget_min_pixels=int(total_min_pixels), budget_max_pixels=int(total_max_pixels), ) + result[VIDEO_BUDGET_RESIZED_KEY] = bool(processor.use_smart_resize) if timing is not None: result[SOURCE_VIDEO_TIMING_KEY] = timing if uses_source_timestamps: diff --git a/cosmos_framework/data/generator/sample_token_counter.py b/cosmos_framework/data/generator/sample_token_counter.py new file mode 100644 index 000000000..fa8e77f88 --- /dev/null +++ b/cosmos_framework/data/generator/sample_token_counter.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""One token accounting implementation for worker mixtures and the parent packer.""" + +import math +from collections.abc import Mapping, Sequence +from typing import Any + +import torch + +from cosmos_framework.model.generator.tokenizers.uniae.frame_math import ( + get_uniae_chunk_frames, + get_uniae_latent_num_frames, + normalize_uniae_chunk_frames, +) +from cosmos_framework.utils.generator.data_utils import read_positive_int_metadata +from cosmos_framework.utils.generator.spatial_patch import normalize_spatial_patch_hw + + +class SampleTokenCounter: + """Price decoded media, text, references and controls before sequence packing.""" + + tokenizer_spatial_compression_factor: int + tokenizer_temporal_compression_factor: int + patch_spatial: int + lidar_spatial_compression: tuple[int, ...] | None + lidar_temporal_compression_factor: int | None + radar_spatial_compression: tuple[int, ...] | None + radar_temporal_compression_factor: int | None + lidar_patch_spatial_hw: tuple[int, int] + radar_patch_spatial_hw: tuple[int, int] + sound_latent_fps: float + audio_sample_rate: int + uniae_pad_frames: int | None + uniae_chunk_frames: int | dict[str, int] | None + + def __init__( + self, + tokenizer_spatial_compression_factor: int, + tokenizer_temporal_compression_factor: int, + patch_spatial: int, + lidar_spatial_compression: Sequence[int] | None = None, + lidar_temporal_compression_factor: int | None = None, + radar_spatial_compression: Sequence[int] | None = None, + radar_temporal_compression_factor: int | None = None, + sound_latent_fps: float = 0, + audio_sample_rate: int = 48000, + uniae_chunk_frames: int | Mapping[str, int] | None = None, + uniae_pad_frames: int | None = None, + lidar_patch_spatial_hw: int | tuple[int, int] | None = None, + radar_patch_spatial_hw: int | tuple[int, int] | None = None, + ) -> None: + self.tokenizer_spatial_compression_factor = tokenizer_spatial_compression_factor + self.tokenizer_temporal_compression_factor = tokenizer_temporal_compression_factor + self.lidar_spatial_compression = ( + tuple(int(factor) for factor in lidar_spatial_compression) + if lidar_spatial_compression is not None + else None + ) + if self.lidar_spatial_compression is not None and ( + len(self.lidar_spatial_compression) != 2 or any(factor <= 0 for factor in self.lidar_spatial_compression) + ): + raise ValueError( + "lidar_spatial_compression must contain two positive factors " + f"(height, width), got {self.lidar_spatial_compression}" + ) + self.lidar_temporal_compression_factor = ( + int(lidar_temporal_compression_factor) if lidar_temporal_compression_factor is not None else None + ) + if self.lidar_temporal_compression_factor is not None and self.lidar_temporal_compression_factor <= 0: + raise ValueError( + f"lidar_temporal_compression_factor must be positive, got {self.lidar_temporal_compression_factor}" + ) + self.radar_spatial_compression = ( + tuple(int(factor) for factor in radar_spatial_compression) + if radar_spatial_compression is not None + else None + ) + if self.radar_spatial_compression is not None and ( + len(self.radar_spatial_compression) != 2 or any(factor <= 0 for factor in self.radar_spatial_compression) + ): + raise ValueError( + "radar_spatial_compression must contain two positive factors " + f"(height, width), got {self.radar_spatial_compression}" + ) + self.radar_temporal_compression_factor = ( + int(radar_temporal_compression_factor) if radar_temporal_compression_factor is not None else None + ) + if self.radar_temporal_compression_factor is not None and self.radar_temporal_compression_factor <= 0: + raise ValueError( + f"radar_temporal_compression_factor must be positive, got {self.radar_temporal_compression_factor}" + ) + self.patch_spatial = patch_spatial + self.lidar_patch_spatial_hw: tuple[int, int] = normalize_spatial_patch_hw( + patch_spatial if lidar_patch_spatial_hw is None else lidar_patch_spatial_hw + ) + self.radar_patch_spatial_hw: tuple[int, int] = normalize_spatial_patch_hw( + patch_spatial if radar_patch_spatial_hw is None else radar_patch_spatial_hw + ) + self.sound_latent_fps = sound_latent_fps + self.audio_sample_rate = audio_sample_rate + self.uniae_pad_frames = int(uniae_pad_frames) if uniae_pad_frames is not None else None + self.uniae_chunk_frames = self._normalize_uniae_chunk_frames(uniae_chunk_frames) + + def count_raw_sample(self, sample: Mapping[str, Any]) -> int: + """Match the packer's single-sample representation without copying tensors. + + Collation and the parent's sample splitter wrap bare text/action/sound + tensors in lists. Media already accept either bare tensors or lists. + """ + data = dict(sample) + for key in ("text_token_ids", "action", "sound"): + value = data.get(key) + if value is not None and not isinstance(value, list): + data[key] = [value] + return self._compute_num_tokens_per_sample(data) + + def _num_lidar_tokens(self, data_batch: Mapping[str, Any]) -> int: + """Cost the sample's LiDAR range clips with the LiDAR VAE's own compression. + + The LiDAR VAE does not compress time, so costing a sweep clip with the camera's 4x + would undercount it fourfold and silently over-pack the batch. + """ + clips = data_batch.get("lidar") + if not clips: + return 0 + if self.lidar_spatial_compression is None or self.lidar_temporal_compression_factor is None: + raise ValueError( + "This batch carries a LiDAR stream, but the loader has no LiDAR compression factors. " + "Set lidar_spatial_compression and lidar_temporal_compression_factor." + ) + spatial_h, spatial_w = self.lidar_spatial_compression + num_tokens = 0 + for clip in clips: + _, T, H, W = clip.shape + patch_h = math.ceil(H // spatial_h / self.lidar_patch_spatial_hw[0]) + patch_w = math.ceil(W // spatial_w / self.lidar_patch_spatial_hw[1]) + latent_t = 1 + (T - 1) // self.lidar_temporal_compression_factor + num_tokens += patch_h * patch_w * latent_t + return num_tokens + + def _num_radar_tokens(self, data_batch: Mapping[str, Any]) -> int: + """Cost the sample's radar BEV clips with the radar VAE's own compression. + + Same arithmetic as the sweeps over a square BEV grid rather than a range image: the + radar VAE does not compress time either, so one scan is one latent frame and the + camera's 4x would undercount a clip fourfold. Radar is the more expensive of the two + sensors per second of clip -- it cycles at ~20 Hz against LiDAR's 10 and each scan + carries more latents -- so an underpriced clip over-packs the batch by more here. + """ + clips = data_batch.get("radar") + if not clips: + return 0 + if self.radar_spatial_compression is None or self.radar_temporal_compression_factor is None: + raise ValueError( + "This batch carries a radar stream, but the loader has no radar compression factors. " + "Set radar_spatial_compression and radar_temporal_compression_factor." + ) + spatial_h, spatial_w = self.radar_spatial_compression + num_tokens = 0 + for clip in clips: + _, T, H, W = clip.shape + patch_h = math.ceil(H // spatial_h / self.radar_patch_spatial_hw[0]) + patch_w = math.ceil(W // spatial_w / self.radar_patch_spatial_hw[1]) + latent_t = 1 + (T - 1) // self.radar_temporal_compression_factor + num_tokens += patch_h * patch_w * latent_t + return num_tokens + + def _normalize_uniae_chunk_frames( + self, uniae_chunk_frames: int | Mapping[str, int] | None + ) -> int | dict[str, int] | None: + return normalize_uniae_chunk_frames( + uniae_chunk_frames, + pad_frames=self.uniae_pad_frames, + temporal_compression_factor=self.tokenizer_temporal_compression_factor, + temporal_divisibility_name="tokenizer_temporal_compression_factor", + ) + + def _get_uniae_chunk_frames(self, spatial_shape: tuple[int, int]) -> int: + assert self.uniae_chunk_frames is not None + return get_uniae_chunk_frames(self.uniae_chunk_frames, spatial_shape=spatial_shape) + + def _compute_vision_latent_t_shape(self, T: int, H: int, W: int) -> int: + if T < 1: + raise ValueError(f"Vision media must contain at least one frame, got {T}.") + if T == 1 or self.uniae_chunk_frames is None: + return 1 + (T - 1) // self.tokenizer_temporal_compression_factor + + assert self.uniae_pad_frames is not None + return get_uniae_latent_num_frames( + T, + self.uniae_chunk_frames, + pad_frames=self.uniae_pad_frames, + temporal_compression_factor=self.tokenizer_temporal_compression_factor, + spatial_shape=(H, W), + ) + + def _compute_num_tokens_per_sample(self, data_batch: Mapping[str, Any]) -> int: + """ + This function computes the number of tokens per sample in the data batch. + This includes text + vision generation tokens + action tokens. + + Args: + data_batch (dict): The data batch containing the text tokens. + + Returns: + int: The number of tokens per sample. + """ + und_tokens, gen_tokens = self._compute_token_split_per_sample(data_batch) + return und_tokens + gen_tokens + + def _compute_token_split_per_sample(self, data_batch: Mapping[str, Any]) -> tuple[int, int]: + """ + Split one sample's packed length into the understanding and generation pathways. + + The two towers of the MoT attend over the same packed sequence but cost + different amounts per token, so the cost model prices them separately. The + two returned counts always sum to ``_compute_num_tokens_per_sample``. + + The three structural markers (````, ````, ````) + are attributed to the stream they delimit rather than broken out, which keeps + that sum exact at the cost of a handful of tokens landing on one side or the + other. + + Args: + data_batch (dict): The data batch containing the text tokens. + + Returns: + tuple[int, int]: ``(understanding tokens, generation tokens)``. + """ + + # The token sequence we have is + # [] + # The spatial dimension of image tokens is compressed by + # vae spatial downsampling factor + pathification + # The temporal dimension of image tokens is compressed by + # vae temporal downsampling factor + # Action tokens have 1 token per time step (no spatial dimension) + + has_text_tokens = "text_token_ids" in data_batch + und_tokens = 0 + gen_tokens = 0 + if has_text_tokens: + text_token_ids = data_batch["text_token_ids"] + if isinstance(text_token_ids, list): + # Price every independently tokenized per-view caption. The packing MR + # will use these boundaries as separate UND segments; accounting for the + # wrapper per segment here prevents the batch builder from underpacking. + und_tokens = sum(int(tokens.shape[0]) + 1 for tokens in text_token_ids) + else: + num_text_tokens = text_token_ids.shape[1] + und_tokens = num_text_tokens + 1 + + # Vision part. A LiDAR-only sample has no camera stream, so the loop below prices + # nothing and the whole generation cost comes from the sweeps. + is_image_batch = "images" in data_batch + # Absent rather than empty for a LiDAR-only sample, and tested for explicitly: the loop + # below also accepts a bare tensor, and `or []` would ask that tensor for its truth + # value before reaching that branch. + input_images_or_videos = data_batch.get("images" if is_image_batch else "video") + if input_images_or_videos is None: + input_images_or_videos = [] + if "enable_per_camera_vae_encoding" in data_batch: + sample_n_views_values = read_positive_int_metadata(data_batch, "sample_n_views", expected_count=1) + frames_per_view_values = read_positive_int_metadata( + data_batch, + "num_video_frames_per_view", + expected_count=1, + ) + if sample_n_views_values is None or frames_per_view_values is None: + raise ValueError( + "sample_n_views and num_video_frames_per_view are required when per-camera VAE encoding is enabled." + ) + sample_n_views = sample_n_views_values[0] + frames_per_view = frames_per_view_values[0] + else: + sample_n_views = None + frames_per_view = None + + # iterate over all the media in the batch + for media in input_images_or_videos if isinstance(input_images_or_videos, list) else [input_images_or_videos]: + if is_image_batch: + _, H, W = media.shape + T = 1 + else: + _, T, H, W = media.shape + + latent_h_shape = H // self.tokenizer_spatial_compression_factor + latent_w_shape = W // self.tokenizer_spatial_compression_factor + patch_h_shape = math.ceil(latent_h_shape / self.patch_spatial) + patch_w_shape = math.ceil(latent_w_shape / self.patch_spatial) + if sample_n_views is not None and frames_per_view is not None and not is_image_batch: + # The multiview dataset resizes every camera to the same H/W and + # selects the same synchronized frame count before concatenating + # the camera-major clips along T. + expected_frames = sample_n_views * frames_per_view + if T != expected_frames: + raise ValueError( + "Multiview vision length must equal sample_n_views * num_video_frames_per_view: " + f"got T={T}, sample_n_views={sample_n_views}, num_video_frames_per_view={frames_per_view}." + ) + latent_t_shape = sample_n_views * self._compute_vision_latent_t_shape(frames_per_view, H, W) + else: + latent_t_shape = self._compute_vision_latent_t_shape(T, H, W) + + num_vision_tokens = patch_h_shape * patch_w_shape * latent_t_shape + if has_text_tokens: + num_vision_tokens += 2 + gen_tokens += num_vision_tokens + + # LiDAR part: its own VAE, hence its own compression factors. Charged to + # gen_tokens so the iteration-time cost model sees the sweeps. + gen_tokens += self._num_lidar_tokens(data_batch) + + # Radar part: a third VAE and a third pair of factors, charged the same way. A joint + # camera + radar sample carries no sweeps and a joint camera + LiDAR one no scans, so + # each of the two calls prices nothing on the other's recipe. + gen_tokens += self._num_radar_tokens(data_batch) + + # Action part: each action time step is 1 token. + # Action tensor shape is (T_action, D) per sample; stored as a single-element list. + if "action" in data_batch: + list_of_actions = data_batch["action"] + for action in list_of_actions: + # skip None actions + if action is None: + continue + num_action_tokens = action.shape[0] + gen_tokens += num_action_tokens + + # Sound part — estimate sound tokens from audio waveform length + if self.sound_latent_fps > 0 and "sound" in data_batch: + sound_data = data_batch["sound"] + if isinstance(sound_data, list) and len(sound_data) > 0: + first_sound = sound_data[0] + # Unwrap nested list if needed + if isinstance(first_sound, list): + first_sound = first_sound[0] + if first_sound is not None and isinstance(first_sound, torch.Tensor): + num_audio_samples = first_sound.shape[-1] + audio_duration = num_audio_samples / self.audio_sample_rate + num_sound_tokens = int(audio_duration * self.sound_latent_fps) + gen_tokens += num_sound_tokens + + return und_tokens, gen_tokens diff --git a/cosmos_framework/data/generator/sequence_packing/autoregressive.py b/cosmos_framework/data/generator/sequence_packing/autoregressive.py index e574f1fc7..6b4a86e24 100644 --- a/cosmos_framework/data/generator/sequence_packing/autoregressive.py +++ b/cosmos_framework/data/generator/sequence_packing/autoregressive.py @@ -195,6 +195,7 @@ def pack_input_sequence_autoregressive( vision_temporal_positions: torch.Tensor | None = None, num_views: int | None = None, text_view_ids: list[int] | None = None, + vision_view_ids: torch.Tensor | None = None, # [V] physical camera IDs, or [-1] for non-rig RGB ) -> PackedSequence: """ Pack input sequence for autoregressive video generation (one AR unit at a time). @@ -267,6 +268,8 @@ def pack_input_sequence_autoregressive( axis. Required by multiview FlexAttention metadata. text_view_ids: Camera-view ID for every entry in a nested ``text_tokens`` payload. ``None`` keeps one sample-level caption visible to all views. + vision_view_ids: Physical RGB camera IDs, independent of caption-local + view IDs. A singleton -1 denotes non-rig video. Returns: Finalized PackedSequence containing the supertoken(s) for this AR unit @@ -415,6 +418,7 @@ def pack_input_sequence_autoregressive( action_domain_id=action_domain_id_list, raw_action_dim=raw_action_dim_list, num_views_per_vision_item=[num_views] if num_views is not None else None, + vision_view_ids=[vision_view_ids] if vision_view_ids is not None else None, ) # Prepare text indexes diff --git a/cosmos_framework/data/generator/sequence_packing/packers.py b/cosmos_framework/data/generator/sequence_packing/packers.py index eb5602eda..957f36c77 100644 --- a/cosmos_framework/data/generator/sequence_packing/packers.py +++ b/cosmos_framework/data/generator/sequence_packing/packers.py @@ -318,7 +318,9 @@ def pack_multiview_action_conditioning( num_views=num_views, ) builder.finish_sample(builder.current_seq_index - generation_start, builder.current_seq_index - sample_start) - return builder.finalize(gen_data_clean=data) + packed = builder.finalize(gen_data_clean=data) + packed.multiview_action_temporal_factor = temporal_factor + return packed def pack_input_sequence( diff --git a/cosmos_framework/data/generator/sequence_packing/packers_test.py b/cosmos_framework/data/generator/sequence_packing/packers_test.py index 99f1534c5..47ce81119 100644 --- a/cosmos_framework/data/generator/sequence_packing/packers_test.py +++ b/cosmos_framework/data/generator/sequence_packing/packers_test.py @@ -16,6 +16,32 @@ from cosmos_framework.data.generator.sequence_packing.sequence import SequencePlan +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("view_id", [-1, 8]) +@pytest.mark.parametrize("frame_idx", [0, 2]) +def test_autoregressive_pack_preserves_physical_view_identity(view_id: int, frame_idx: int) -> None: + from cosmos_framework.data.generator.sequence_packing.autoregressive import pack_input_sequence_autoregressive + + ids = torch.tensor([view_id]) # [V] + packed = pack_input_sequence_autoregressive( + vision_latent=torch.zeros(1, 4, 1, 2, 2), # [B,C,T,H,W] + action_latent=None, + text_tokens=[1, 2] if frame_idx == 0 else None, + timestep=0.0, + fps_vision=[24.0], + fps_action=None, + special_tokens={"eos_token_id": 3, "start_of_generation": 4, "end_of_generation": 5}, + frame_idx=frame_idx, + cached_text_offset=None if frame_idx == 0 else 4, + vision_view_ids=ids, + ) + + assert packed.vision_view_ids is not None + assert len(packed.vision_view_ids) == 1 + torch.testing.assert_close(packed.vision_view_ids[0], ids) + + @pytest.mark.L0 @pytest.mark.CPU @pytest.mark.parametrize("fps", [None, 0.0, -1.0, float("nan"), float("inf"), -float("inf")]) diff --git a/cosmos_framework/data/generator/sequence_packing/sequence.py b/cosmos_framework/data/generator/sequence_packing/sequence.py index f83874c25..d22e50e0a 100644 --- a/cosmos_framework/data/generator/sequence_packing/sequence.py +++ b/cosmos_framework/data/generator/sequence_packing/sequence.py @@ -1383,6 +1383,12 @@ class PackedSequence: # with samples that have none. num_action_items_per_sample: list[int] | None = None + # Separate-camera RGB/action layout: each camera is one vision item, and its + # full-rate actions occupy this many rows per latent step, starting at step 1. + # Set by the dedicated packer so attention reads the layout from the batch, + # without interpreting arbitrary control/target items as separate cameras. + multiview_action_temporal_factor: int | None = None + def __post_init__(self) -> None: self._sequence_pack_metadata: SequencePackMetadata | None = None assert isinstance(self.text_ids, torch.Tensor), "PackedSequence.text_ids must be finalized" diff --git a/cosmos_framework/data/generator/token_mix_control.py b/cosmos_framework/data/generator/token_mix_control.py index 06d1828f6..110793860 100644 --- a/cosmos_framework/data/generator/token_mix_control.py +++ b/cosmos_framework/data/generator/token_mix_control.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -"""Holds the realized token mix on the configured ratios when packing is cost-based. +"""Adapt source draw probabilities toward the configured token mix. A loader that draws one source per batch spends the whole batch on that source, so over ``N`` batches source ``i`` contributes ``N * p_i * T_i`` tokens, where ``p_i`` is @@ -21,16 +21,25 @@ cooperation from the packer: it prices what the packer did, drops, look-ahead misses, token ceiling and all. -The same estimator supports sample-level mixing: each observation is one fresh, -packable candidate and ``T_i`` is mean tokens per sample. Deferred candidates must -be observed only once. Their eventual acceptance is accounted for by the loader, -separately from the candidate estimate, using ``record_realized=False`` here. - -This is a plug-in rather than a servo. ``T_i`` is a property of a source's data under -the packer's ceilings and does not depend on how often that source is drawn, so the -loop has no feedback path into the quantity it measures: no gain to tune, nothing to -ring, and the mix lands on target as soon as the EMA has warmed up. What is left is -estimator noise, traded against tracking speed by ``half_life_batches``. +The same estimator supports sample-level mixing, where ``T_i`` is mean tokens per +sample. ``RandomMixedJointDataLoader`` observes each sample only when accepted into +a packed batch. Deferred candidates do not update the estimate or warmup until +accepted. ``TokenMixedDataset`` instead observes packable samples before yielding +them to its parent loader. Packing statistics are reported by loader callbacks +rather than this controller. + +The default correction estimates ``T_i``, a property of a source's data under the +packer's ceilings that does not depend on how often that source is drawn. It has no +feedback path into the quantity it measures. Estimator noise is traded against +tracking speed by ``half_life_batches``. + +The mixed loader also enables best-effort feedback: multiply the size correction +by target token share / observed token share since this controller was created. +Overrepresented sources get smaller weights on subsequent draws. Its observations +count accepted tokens only, so deferral does not charge a source before acceptance. +This is a soft correction: every eligible source keeps a positive weight, +and fallback to the only eligible source still draws it regardless of its surplus. +Thus packing constraints can prevent the realized mix from reaching its target. Every rank runs its own controller over its own batches, which keeps it out of the collective path -- the packer's daemon thread makes any collective there a hazard. @@ -39,8 +48,9 @@ ``IterativeJointDataLoader`` seeds its draw so that every rank picks the same source at a given iteration, and rank-local weights would break that. -The state is not checkpointed. A resumed job re-estimates from its first batches and -is back on target within a few half-lives. +The state is not checkpointed. A resumed job re-estimates from its first batches; +its sample-size estimates settle within a few half-lives and its feedback history +starts over. """ from __future__ import annotations @@ -50,8 +60,6 @@ import numpy as np -from cosmos_framework.utils import log - @dataclass class TokenMixControlConfig: @@ -75,22 +83,25 @@ class TokenMixControlConfig: since one batch of an unlucky shape is a poor estimate of a source. max_weight_ratio: Widest factor by which a draw probability may depart from its configured share, in either direction, and equivalently the widest - spread in tokens per batch the draw will fully invert. A guard rail + spread in tokens per batch the default correction will fully invert. + When feedback is enabled, this bounds the combined correction. A guard rail against a bad estimate starving a source rather than a tuning knob, but one with a cost: a mixture whose tokens per batch spread wider than this is only partly corrected, and the residual shows up as a realized share that sits off target with the source pinned at the bound. ``None`` disables clipping, allowing arbitrary sample-length differences. - log_every_n_batches: How often to log target against realized share. The - realized figure is measured over that window, so it is also how long - the diagnostic takes to reflect a change. + feedback: Apply target share / cumulative observed token share after + warmup, in addition to the sample-size correction. Observed tokens + count accepted samples in ``RandomMixedJointDataLoader``. + This history is independent of the EMA half-life. + It provides soft feedback without excluding overrepresented sources. """ enabled: bool = True half_life_batches: float = 50.0 warmup_batches_per_source: int = 8 max_weight_ratio: float | None = 8.0 - log_every_n_batches: int = 500 + feedback: bool = False def build(self, dataset_names: Sequence[str], target_shares: Sequence[float]) -> TokenMixController | None: """Resolve into a controller, or ``None`` when there is nothing to control. @@ -114,8 +125,6 @@ def build(self, dataset_names: Sequence[str], target_shares: Sequence[float]) -> raise ValueError(f"warmup_batches_per_source must be at least 1, got {self.warmup_batches_per_source}.") if self.max_weight_ratio is not None and self.max_weight_ratio < 1.0: raise ValueError(f"max_weight_ratio must be at least 1, got {self.max_weight_ratio}.") - if self.log_every_n_batches < 1: - raise ValueError(f"log_every_n_batches must be at least 1, got {self.log_every_n_batches}.") if len(dataset_names) != len(target_shares): raise ValueError(f"Got {len(dataset_names)} dataset names and {len(target_shares)} target shares.") @@ -125,12 +134,14 @@ def build(self, dataset_names: Sequence[str], target_shares: Sequence[float]) -> class TokenMixController: - """Draw probabilities that keep each source's token share on its configured one. + """Draw probabilities that target each source's configured token share. Attributes: sampling_probs: The distribution over all sources after the last observation. Use ``probabilities(active)`` when sources can exhaust. draw_counts: Number of positive-token observations from each source. + observed_tokens: Cumulative token exposure, counting each observation + once. ``RandomMixedJointDataLoader`` counts accepted samples only. """ def __init__( @@ -154,25 +165,21 @@ def __init__( self._decay: float = 0.5 ** (1.0 / config.half_life_batches) self._mean_tokens: np.ndarray = np.zeros(len(self.names), dtype=np.float64) self.draw_counts: np.ndarray = np.zeros(len(self.names), dtype=np.int64) - self._window_tokens: np.ndarray = np.zeros(len(self.names), dtype=np.float64) - self._window_batches: int = 0 - self._at_band: np.ndarray = np.zeros(len(self.names), dtype=bool) + self.observed_tokens: np.ndarray = np.zeros(len(self.names), dtype=np.int64) - def observe(self, index: int, tokens: int, *, record_realized: bool = True) -> None: + def observe(self, index: int, tokens: int) -> None: """Fold one source draw into the estimate and refresh the probabilities. Args: index: Source of this observation, in the loader's index order. tokens: Tokens in the batch or candidate. Non-positive counts are ignored, an empty draw being a fact about the input stream. - record_realized: Record this observation in the accepted-token window. - Batch draws are already packed. Sample candidates may be deferred, - so their loader passes false and tracks accepted tokens separately. """ if tokens <= 0: return self.draw_counts[index] += 1 + self.observed_tokens[index] += tokens if self.draw_counts[index] == 1: # Seed rather than decay from zero, which would read a source's first # observations as cheap and draw it even harder while the error was largest. @@ -181,75 +188,45 @@ def observe(self, index: int, tokens: int, *, record_realized: bool = True) -> N self._mean_tokens[index] *= self._decay self._mean_tokens[index] += (1.0 - self._decay) * tokens self.sampling_probs = self.probabilities() - if not record_realized: - return - - self._window_tokens[index] += tokens - self._window_batches += 1 - - if self._window_batches >= self.config.log_every_n_batches: - # Rank 0 only, for volume: every rank holds its own controller, but they see - # the same data and differ by estimator noise, so one rank's view is the mix. - log.info(self.describe()) - self._window_tokens[:] = 0.0 - self._window_batches = 0 - - def realized_shares(self) -> np.ndarray: - """Token share each source won since the last log, or zeros before any batch.""" - total = self._window_tokens.sum() - if total <= 0.0: - return np.zeros_like(self._window_tokens) - return self._window_tokens / total - - def describe(self) -> str: - """Target against realized share and the draw in force, one line per source.""" - realized = self.realized_shares() - warming = int((self.draw_counts < self.config.warmup_batches_per_source).sum()) - header = f"token mix over {self._window_batches} batches" - if warming: - header += f" (warming up: {warming} source(s) below {self.config.warmup_batches_per_source} observations)" - lines = [header] - for index, name in enumerate(self.names): - # Naming the sources at the band is the whole diagnostic for a mix that is - # off target and staying there: it says to widen max_weight_ratio rather - # than to distrust the ratios. - at_band = " (at the band, so only partly corrected)" if self._at_band[index] else "" - lines.append( - f" {name}: target {self.target_shares[index]:.1%}, realized {realized[index]:.1%}, " - f"drawn at {self.sampling_probs[index]:.1%}, {self._mean_tokens[index]:,.0f} tokens/observation{at_band}" - ) - return "\n".join(lines) def probabilities(self, active: np.ndarray | None = None) -> np.ndarray: """``ratio_i / T_i``, normalised over the active sources after warmup. + With feedback, multiply by target share / cumulative observed share. + Shares are normalized over the same active sources, so exhausted or + temporarily excluded sources do not change their relative correction. + Warmup guarantees positive exposure for every corrected source. + Disabled controllers return the configured ratios over active sources. An all-false mask returns zeros; the caller must stop drawing in that case. Inactive sources affect neither warmup nor correction clipping. - When configured, ``1 / T_i`` is clipped in log space around its geometric mean - rather than after normalising. Ratios between corrections survive - normalisation where absolute values do not, so bounding their spread there is + When configured, the combined size and feedback correction is clipped in + log space around its geometric mean rather than after normalising. Ratios + between corrections survive normalisation where absolute values do not, so bounding their spread there is what makes ``max_weight_ratio`` an exact bound on the departure of every ``p_i`` from its configured share, instead of one that normalising undoes. - - Also records which sources the clip caught, for :meth:`describe`. """ if active is None: active = np.ones(len(self.names), dtype=bool) weights = np.where(active, self.target_shares, 0.0) - self._at_band[:] = False if not active.any(): return weights if not self.config.enabled or (self.draw_counts[active] < self.config.warmup_batches_per_source).any(): return weights / weights.sum() + correction = 1.0 / self._mean_tokens[active] + if self.config.feedback: + exposure = self.observed_tokens[active] + observed_shares = exposure / exposure.sum() + target_shares = weights[active] / weights.sum() + correction *= target_shares / observed_shares + if self.config.max_weight_ratio is None: - weights[active] /= self._mean_tokens[active] + weights[active] *= correction else: - correction = -np.log(self._mean_tokens[active]) + correction = np.log(correction) correction -= correction.mean() half_band = 0.5 * np.log(self.config.max_weight_ratio) - self._at_band[active] = np.abs(correction) > half_band weights[active] *= np.exp(np.clip(correction, -half_band, half_band)) return weights / weights.sum() diff --git a/cosmos_framework/data/generator/transfer_length_index.py b/cosmos_framework/data/generator/transfer_length_index.py index c4ff98692..228243c4d 100644 --- a/cosmos_framework/data/generator/transfer_length_index.py +++ b/cosmos_framework/data/generator/transfer_length_index.py @@ -3,13 +3,13 @@ """Build a complete, version-pinned Transfer video census and disjoint length index. -Only canonical Lance metadata is read: no video bytes are downloaded. The runtime +Only source metadata rows are read: no video bytes are downloaded. The runtime adapter supplies candidate selection so source/caption/control eligibility cannot silently diverge. Every unique RGB URI belongs to one bucket, with its longest fully available caption window. The full window is resampled, never prefix-cropped. A manifest is published only after scanning every source row and closing SQLite. -Run with ``python -m cosmos_framework.data.generator.transfer_length_index``. +Build an index from metadata rows with ``build_index_from_rows``. """ from __future__ import annotations diff --git a/cosmos_framework/inference/common/distillation_export_test.py b/cosmos_framework/inference/common/distillation_export_test.py index 697f4a281..3bcd50360 100644 --- a/cosmos_framework/inference/common/distillation_export_test.py +++ b/cosmos_framework/inference/common/distillation_export_test.py @@ -5,6 +5,11 @@ import attrs import pytest +from cosmos3.common import public_model_config +from cosmos3.common.public_model_config import ( + build_public_model_config, + restore_model_config_from_public_model_config, +) from cosmos_framework.inference.common import distillation_export from cosmos_framework.inference.common.distillation_export import ( @@ -321,6 +326,49 @@ def test_sanitize_student_public_model_config_preserves_qwen_variant(model_size: ) +def test_sanitized_student_config_round_trips_through_public_aliases(monkeypatch: pytest.MonkeyPatch) -> None: + model_dict = { + "_target_": "cosmos_framework.model.generator.omni_mot_model.OmniMoTModel", + "config": { + "_type": "cosmos_framework.configs.base.defaults.model_config.OmniMoTModelConfig", + "vlm_config": { + "_type": "cosmos_framework.configs.base.defaults.reasoner.VLMConfig", + "model_instance": { + "_target_": "cosmos_framework.model.generator.mot.unified_mot.Qwen3VLTextForCausalLM", + "config": { + "_target_": "cosmos_framework.configs.base.defaults.reasoner.create_vlm_config", + "base_config": { + "_target_": ( + "cosmos_framework.model.generator.mot.unified_mot.Qwen3VLMoTConfig.from_json_file" + ), + "json_file": "/private/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json", + }, + }, + }, + }, + }, + } + + distillation_export.sanitize_student_public_model_config(model_dict) + exported_model_config = build_public_model_config(model_dict) + + exported_base_config = exported_model_config["config"]["vlm_config"]["model_instance"]["config"]["base_config"] + assert exported_model_config["_target"] == "omni_mot_model" + assert exported_base_config["json_file"] == ("cosmos3://vfm/models/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json") + + monkeypatch.setattr( + public_model_config, + "_module_exists", + lambda module: module == "projects.cosmos3.cosmos3", + ) + restored_model_config = restore_model_config_from_public_model_config(exported_model_config) + + restored_base_config = restored_model_config["config"]["vlm_config"]["model_instance"]["config"]["base_config"] + assert restored_base_config["json_file"] == ( + "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" + ) + + def test_sanitize_student_public_model_config_is_idempotent() -> None: model_dict = { "config": { diff --git a/cosmos_framework/model/_base.py b/cosmos_framework/model/_base.py index ed9e74dc1..ad7ea08d4 100644 --- a/cosmos_framework/model/_base.py +++ b/cosmos_framework/model/_base.py @@ -20,6 +20,13 @@ class ImaginaireModel(torch.nn.Module): - init_optimizer_scheduler(): Creates the optimizer and scheduler for the model. """ + # Only objectives that need future CPU labels opt into trainer lookahead. + requires_accumulation_window: bool = False + + def prepare_accumulation_window(self, batches: tuple[dict[str, Any], ...]) -> None: + """Inspect one complete CPU accumulation window before its first forward.""" + del batches + def __init__(self) -> None: super().__init__() self.parallel_dims = None diff --git a/cosmos_framework/model/generator/algorithm/loss/cross_entropy.py b/cosmos_framework/model/generator/algorithm/loss/cross_entropy.py index 18fd3ffde..4d3de912e 100644 --- a/cosmos_framework/model/generator/algorithm/loss/cross_entropy.py +++ b/cosmos_framework/model/generator/algorithm/loss/cross_entropy.py @@ -22,27 +22,30 @@ Both losses reduce their normalizer (valid-token count for plain CE, the exponent-weighted sample weight for weighted CE) over the WHOLE WORLD, taking no group argument. Every rank must hold a different sample for that to be the count the -formula wants, which holds for VLM: ``ParallelDims`` pins -``dp_replicate * dp_shard == world_size``, so the data-parallel mesh is the world, and -the axes that would put the same sample on several ranks are rejected — cp/cfgp by -``VLMModel``, tp/pp by not existing. Revisit this if any of that changes. +formula wants on the logits-based path: ``ParallelDims`` pins +``dp_replicate * dp_shard == world_size``, so the data-parallel mesh is the world. +CP is accepted only by the separately normalized fused path described below; +CFGP remains unsupported, and TP/PP are not implemented here. async_safe_ce instead reduces over the dp_shard sub-group, which under HSDP normalizes each replicate group separately; that agrees with the global token mean only when those groups carry equal token counts. -Neither loss takes a cp_group: reducing across context-parallel segments is out of -scope here. +The two logits-based losses remain CP-disabled. The fused weighted loss additionally +accepts a context_parallel_group and globally shifted, locally sliced labels; its +per-sample counts are reduced over CP before applying the sample weighting. """ from __future__ import annotations from dataclasses import dataclass +from typing import Any import torch import torch.distributed as dist import torch.nn.functional as F from torch.distributed.tensor import DTensor +from torch.utils.checkpoint import checkpoint from cosmos_framework.utils.generator.input_probe import maybe_dump_loss_reduction from cosmos_framework.utils.generator.reasoner.constant import IGNORE_INDEX @@ -231,6 +234,44 @@ def weighted_cross_entropy_loss( return loss +def weighted_ce_denominator_from_labels( + labels: torch.Tensor, + *, + exponent: float, + cu_seq_lens: torch.Tensor | None = None, + ignore_index: int = IGNORE_INDEX, +) -> torch.Tensor: + """Count the weighted objective mass from full, unsharded CPU labels. + + This is label-only lookahead, before CP shifts/slices the sequence. Packed + boundaries retain per-conversation weighting; ignored-only segments add zero. + The caller handles CP replication and the WORLD reduction. + """ + if labels.device.type != "cpu" or labels.ndim != 2 or labels.shape[1] < 2: + raise ValueError("Accumulation lookahead requires full CPU labels [B,T], T >= 2") + if not 0.0 <= exponent <= 1.0: + raise ValueError("Weighted CE exponent must be in [0,1]") + valid = labels != ignore_index + valid[:, 0] = False # Match the global next-token shift, without changing labels. + if cu_seq_lens is None: + counts = valid.sum(dim=1).float() + else: + if ( + labels.shape[0] != 1 + or cu_seq_lens.device.type != "cpu" + or cu_seq_lens.dtype not in (torch.int32, torch.int64) + or cu_seq_lens.ndim != 1 + or cu_seq_lens.numel() < 2 + or cu_seq_lens[0] != 0 + or cu_seq_lens[-1] != labels.shape[1] + or not bool((cu_seq_lens[1:] > cu_seq_lens[:-1]).all()) + ): + raise ValueError("Packed lookahead boundaries must cover one complete CPU row") + prefix = torch.cat((torch.zeros(1, dtype=torch.long), valid[0].cumsum(0))) + counts = (prefix[cu_seq_lens[1:]] - prefix[cu_seq_lens[:-1]]).float() + return torch.where(counts > 0, counts.clamp(min=1).pow(1 - exponent), 0.0).sum() + + def _fused_linear_ce_sum( hidden: torch.Tensor, weight: torch.Tensor, labels: torch.Tensor, ignore_index: int ) -> torch.Tensor: @@ -242,6 +283,42 @@ def _fused_linear_ce_sum( ) +class _FP32GradientAnchor(torch.autograd.Function): + """Carry a matrix-shaped FP32 gradient without copying the forward weight. + + Only shape, dtype and the autograd edge are used by the compute-weight + aliases below. The expanded scalar's values never enter the loss kernel. + """ + + @staticmethod + def forward(ctx: Any, weight: torch.Tensor) -> torch.Tensor: + ctx.weight_dtype = weight.dtype + return weight.new_zeros((), dtype=torch.float32).expand_as(weight) + + @staticmethod + def backward(ctx: Any, gradient: torch.Tensor) -> torch.Tensor: + return gradient.to(ctx.weight_dtype) + + +class _ComputeWeightWithFP32Gradient(torch.autograd.Function): + """Alias compute storage while returning each head contribution to one FP32 node. + + All conversations share the original compute weight. A distinct autograd edge + per conversation promotes its gradient before the shared FP32 fan-in; a + single cast then returns the sum to the original gathered weight's dtype. + """ + + @staticmethod + def forward(ctx: Any, accumulation_weight: torch.Tensor, compute_weight: torch.Tensor) -> torch.Tensor: + del ctx, accumulation_weight + return compute_weight + + @staticmethod + def backward(ctx: Any, gradient: torch.Tensor) -> tuple[torch.Tensor, None]: + del ctx + return gradient.float(), None + + def fused_weighted_cross_entropy_loss( hidden_states: torch.Tensor, labels: torch.Tensor, @@ -249,12 +326,24 @@ def fused_weighted_cross_entropy_loss( *, exponent: float, ignore_index: int = IGNORE_INDEX, + context_parallel_group: dist.ProcessGroup | None = None, + cu_seq_lens: torch.Tensor | None = None, + recompute: bool = False, + accumulate_head_grad_fp32: bool = False, ) -> tuple[torch.Tensor, LossStatistics]: """Compute the unified weighted objective without allocating [B,T,vocabulary] logits. Called inside the FSDP model forward, while the lm_head weight is gathered. Each sample uses a scalar fused CE call. The denominator spans the unified - trainer's whole DP world, including HSDP replicas. + trainer's whole DP world, including HSDP replicas. With CP, labels have + already been shifted globally and sliced with the inputs. + Sample weights use CP-global valid counts; each shard contributes 1/CP of + the sample's denominator so FSDP's WORLD average retains the same objective. + Optional recomputation avoids keeping Liger's dense weight-gradient buffer + for every conversation alive until backward; the objective is unchanged. + Optional FP32 accumulation adds the native per-conversation head gradients + before casting once to the gathered weight dtype. CE compute and hidden + gradients keep their original precision, and compute-weight storage is shared. """ if hidden_states.ndim != 3 or labels.shape != hidden_states.shape[:2]: raise ValueError("Fused CE requires aligned hidden states [B,T,H] and labels [B,T]") @@ -265,21 +354,69 @@ def fused_weighted_cross_entropy_loss( if isinstance(lm_head_weight, DTensor): raise ValueError("Fused CE must execute inside the FSDP root with its gathered lm_head weight") weight = lm_head_weight.to(hidden_states.dtype) - targets = labels[:, 1:].contiguous() - counts = (targets != ignore_index).sum(dim=1).to(torch.float32) + targets = labels[:, 1:].contiguous() if context_parallel_group is None else labels.contiguous() + hidden = hidden_states[:, :-1] if context_parallel_group is None else hidden_states + if cu_seq_lens is None: + segments = [(hidden[index], targets[index]) for index in range(labels.shape[0])] + else: + if labels.shape[0] != 1 or cu_seq_lens.ndim != 1 or cu_seq_lens.numel() < 2: + raise ValueError("Packed fused CE requires one row and cumulative conversation boundaries") + boundaries = cu_seq_lens.tolist() + cp_size = 1 if context_parallel_group is None else dist.get_world_size(context_parallel_group) + full_length = labels.shape[1] * cp_size + if ( + boundaries[0] != 0 + or boundaries[-1] != full_length + or any(left >= right for left, right in zip(boundaries[:-1], boundaries[1:], strict=True)) + ): + raise ValueError("Packed boundaries must cover the complete physical sequence") + # Ownership follows the globally shifted target. Conversation-start + # labels are ignored by the pack collator, so no cross-conversation + # prediction contributes to either CE or the per-conversation counts. + target_start = 1 + if context_parallel_group is not None: + target_start += dist.get_rank(context_parallel_group) * labels.shape[1] + segments = [] + for left, right in zip(boundaries[:-1], boundaries[1:], strict=True): + start = min(max(left - target_start, 0), targets.shape[1]) + end = min(max(right - target_start, 0), targets.shape[1]) + segments.append((hidden[0, start:end], targets[0, start:end])) + counts = torch.stack([(target != ignore_index).sum() for _, target in segments]) + sample_counts = counts.detach().to(torch.float32) + cp_size = 1 + if context_parallel_group is not None: + dist.all_reduce(sample_counts, group=context_parallel_group) + cp_size = dist.get_world_size(context_parallel_group) losses = [] - for index in range(labels.shape[0]): + accumulation_weight: torch.Tensor | None = None + promote_head_grad = ( + accumulate_head_grad_fp32 + and torch.is_grad_enabled() + and lm_head_weight.requires_grad + and hidden_states.dtype in (torch.bfloat16, torch.float16) + and len(segments) > 1 + ) + for index, (sample_hidden, sample_targets) in enumerate(segments): if bool(counts[index] > 0): - losses.append( - _fused_linear_ce_sum(hidden_states[index, :-1].contiguous(), weight, targets[index], ignore_index) - ) + sample_weight = weight + if promote_head_grad: + if accumulation_weight is None: + accumulation_weight = _FP32GradientAnchor.apply(lm_head_weight) + sample_weight = _ComputeWeightWithFP32Gradient.apply(accumulation_weight, weight.detach()) + args = (sample_hidden.contiguous(), sample_weight, sample_targets.contiguous(), ignore_index) + if recompute and torch.is_grad_enabled(): + losses.append(checkpoint(_fused_linear_ce_sum, *args, use_reentrant=False)) + else: + losses.append(_fused_linear_ce_sum(*args)) else: # Empty slices keep both branches in autograd without reducing the large weight matrix. - losses.append(hidden_states[index, :0].sum() + weight[:0].sum()) + losses.append(sample_hidden[:0].sum() + weight[:0].sum()) loss_sums = torch.stack(losses).float() - has_valid = counts > 0 - numerator = (loss_sums / counts.clamp(min=1).pow(exponent)).sum() - denominator = torch.where(has_valid, counts.clamp(min=1).pow(1 - exponent), 0.0).sum() + has_valid = sample_counts > 0 + sample_numerators = loss_sums / sample_counts.clamp(min=1).pow(exponent) + sample_denominators = torch.where(has_valid, sample_counts.clamp(min=1).pow(1 - exponent), 0.0) / cp_size + numerator = sample_numerators.sum() + denominator = sample_denominators.sum() global_denominator = denominator.detach().clone() world_size = 1 if dist.is_available() and dist.is_initialized(): @@ -292,8 +429,8 @@ def fused_weighted_cross_entropy_loss( global_objective_denominator=global_denominator.detach(), token_ce_sum=loss_sums.sum().detach(), valid_token_count=counts.sum().detach(), - per_sample_objective_numerator=(loss_sums / counts.clamp(min=1).pow(exponent)).detach(), - per_sample_objective_denominator=torch.where(has_valid, counts.clamp(min=1).pow(1 - exponent), 0.0).detach(), + per_sample_objective_numerator=sample_numerators.detach(), + per_sample_objective_denominator=sample_denominators.detach(), per_sample_token_ce_sum=loss_sums.detach(), per_sample_valid_token_count=counts.detach().to(torch.long), ) diff --git a/cosmos_framework/model/generator/algorithm/loss/flow_matching.py b/cosmos_framework/model/generator/algorithm/loss/flow_matching.py index a2264b8ad..b523dd4b3 100644 --- a/cosmos_framework/model/generator/algorithm/loss/flow_matching.py +++ b/cosmos_framework/model/generator/algorithm/loss/flow_matching.py @@ -100,7 +100,7 @@ def compute_flow_matching_loss( Returns: tuple: A tuple containing two elements: - Flow matching loss (or dummy loss for gradient consistency). - - Per-instance loss (or dummy loss for gradient consistency). + - Per-instance loss, including a zero entry per prediction when fully conditioned. """ if not has_valid_tokens: # Dummy loss to maintain backward graph consistency across ranks @@ -108,7 +108,9 @@ def compute_flow_matching_loss( if item_losses is not None: item_losses.weighted_losses = torch.stack([p.sum() * 0.0 for p in pred]) # [N_items] item_losses.valid = torch.zeros(len(pred), dtype=torch.bool, device=dummy_loss.device) # [N_items] - return dummy_loss, dummy_loss.unsqueeze(0) # make per-instance loss 1-D + # Policy/ID-only packs still have N vision items and N sigmas. Preserve + # that logging contract while keeping the scalar FSDP dummy unchanged. + return dummy_loss, dummy_loss.expand(len(pred)) # [N_items] # condition_mask[i] is T-first with trailing singletons: [T,1,1] vision, [T,1] action. # tw_i gets the same shape so w(σ_t) broadcasts element-wise over non-T dims. diff --git a/cosmos_framework/model/generator/algorithm/loss/fp32_head_gradient_test.py b/cosmos_framework/model/generator/algorithm/loss/fp32_head_gradient_test.py new file mode 100644 index 000000000..96c91c68f --- /dev/null +++ b/cosmos_framework/model/generator/algorithm/loss/fp32_head_gradient_test.py @@ -0,0 +1,139 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Protect small head-gradient contributions across packed conversations.""" + +from dataclasses import fields + +import pytest +import torch +import torch.nn.functional as F + +from cosmos_framework.model.generator.algorithm.loss import cross_entropy as ce + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + + +@pytest.mark.parametrize("recompute", [False, True]) +@pytest.mark.parametrize("weight_dtype", [torch.bfloat16, torch.float32]) +def test_head_accumulation_preserves_small_contributions_and_shares_storage( + monkeypatch: pytest.MonkeyPatch, recompute: bool, weight_dtype: torch.dtype +) -> None: + storage: set[int] = set() + + def scalar_ce(hidden: torch.Tensor, weight: torch.Tensor, labels: torch.Tensor, ignore_index: int) -> torch.Tensor: + storage.add(weight.untyped_storage().data_ptr()) + assert weight.dtype == torch.bfloat16 + return F.cross_entropy(F.linear(hidden, weight).float(), labels, ignore_index=ignore_index, reduction="sum") + + monkeypatch.setattr(ce, "_fused_linear_ce_sum", scalar_ce) + # All four logits are zero, so their derivatives are exactly 1/4 or -3/4. + # The large positive/negative terms cancel; the 62 small terms must survive. + amplitudes = [1024.0] + [1.0] * 31 + [-1024.0] + [1.0] * 31 + hidden = torch.zeros(1, 128, 2, dtype=torch.bfloat16) + hidden[0, ::2, 0] = torch.tensor(amplitudes, dtype=torch.bfloat16) + hidden[0, ::2, 1] = 1 + labels = torch.tensor([[-100, 1] * 64]) + boundaries = torch.arange(0, 129, 2, dtype=torch.int32) + gradients: dict[bool, torch.Tensor] = {} + for promoted in (False, True): + storage.clear() + weight = torch.zeros(4, 2, dtype=weight_dtype, requires_grad=True) + loss, _ = ce.fused_weighted_cross_entropy_loss( + hidden, + labels, + weight, + exponent=0.5, + cu_seq_lens=boundaries, + recompute=recompute, + accumulate_head_grad_fp32=promoted, + ) + loss.backward() + assert weight.grad is not None + gradients[promoted] = weight.grad.detach().clone() + # Recompute may revisit a conversation, but never allocates a head per conversation. + assert len(storage) == 1 + expected = ( + torch.tensor([0.25, -0.75, 0.25, 0.25], dtype=torch.float64)[:, None] + * torch.tensor([62.0 / 64, 1.0], dtype=torch.float64)[None, :] + ) + torch.testing.assert_close(gradients[True], expected.to(weight_dtype), rtol=0, atol=0) + assert not torch.equal(gradients[False], gradients[True]) + + +@pytest.mark.parametrize("recompute", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("all_ignored", [False, True]) +def test_head_accumulation_preserves_forward_statistics_and_hidden_gradient( + monkeypatch: pytest.MonkeyPatch, recompute: bool, dtype: torch.dtype, all_ignored: bool +) -> None: + def scalar_ce(hidden: torch.Tensor, weight: torch.Tensor, labels: torch.Tensor, ignore_index: int) -> torch.Tensor: + return F.cross_entropy(F.linear(hidden, weight).float(), labels, ignore_index=ignore_index, reduction="sum") + + monkeypatch.setattr(ce, "_fused_linear_ce_sum", scalar_ce) + torch.manual_seed(519) + initial_hidden = torch.randn(1, 16, 4, dtype=dtype) + initial_weight = torch.randn(8, 4, dtype=dtype) + labels = torch.tensor([[-100, 1, 2, 3, -100, -100, -100, -100, -100, 2, 3, 1, -100, 6, 5, 7]]) + if all_ignored: + labels.fill_(-100) + boundaries = torch.tensor([0, 4, 8, 12, 16], dtype=torch.int32) + results: list[tuple[torch.Tensor, ce.LossStatistics, torch.Tensor]] = [] + for promoted in (False, True): + hidden = initial_hidden.clone().requires_grad_() + weight = initial_weight.clone().requires_grad_() + loss, statistics = ce.fused_weighted_cross_entropy_loss( + hidden, + labels, + weight, + exponent=0.5, + cu_seq_lens=boundaries, + recompute=recompute, + accumulate_head_grad_fp32=promoted, + ) + (loss * 0.25).backward() + assert hidden.grad is not None and weight.grad is not None + if all_ignored: + assert torch.count_nonzero(weight.grad) == 0 + results.append((loss.detach(), statistics, hidden.grad)) + torch.testing.assert_close(results[0][0], results[1][0], rtol=0, atol=0) + torch.testing.assert_close(results[0][2], results[1][2], rtol=0, atol=0) + for field in fields(ce.LossStatistics): + torch.testing.assert_close( + getattr(results[0][1], field.name), getattr(results[1][1], field.name), rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("grad_enabled", [False, True]) +def test_head_accumulation_does_not_promote_a_frozen_head(monkeypatch: pytest.MonkeyPatch, grad_enabled: bool) -> None: + def scalar_ce(hidden: torch.Tensor, weight: torch.Tensor, labels: torch.Tensor, ignore_index: int) -> torch.Tensor: + assert not weight.requires_grad + return F.cross_entropy(F.linear(hidden, weight).float(), labels, ignore_index=ignore_index, reduction="sum") + + monkeypatch.setattr(ce, "_fused_linear_ce_sum", scalar_ce) + hidden = torch.ones(1, 4, 2, dtype=torch.bfloat16, requires_grad=True) + weight = torch.ones(4, 2, dtype=torch.bfloat16) + with torch.set_grad_enabled(grad_enabled): + loss, _ = ce.fused_weighted_cross_entropy_loss( + hidden, + torch.tensor([[-100, 1, -100, 2]]), + weight, + exponent=0.5, + cu_seq_lens=torch.tensor([0, 2, 4], dtype=torch.int32), + accumulate_head_grad_fp32=True, + ) + assert loss.requires_grad == grad_enabled + if grad_enabled: + loss.backward() + assert hidden.grad is not None + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_fp32_anchor_has_constant_forward_storage_and_preserves_full_gradient(dtype: torch.dtype) -> None: + weight = torch.ones(512, 64, dtype=dtype, requires_grad=True) + anchor = ce._FP32GradientAnchor.apply(weight) + assert anchor.shape == weight.shape and anchor.dtype == torch.float32 + assert anchor.untyped_storage().nbytes() == torch.tensor(0, dtype=torch.float32).element_size() + gradient = torch.linspace(-1, 1, weight.numel(), dtype=torch.float32).reshape_as(weight) + anchor.backward(gradient) + assert weight.grad is not None + torch.testing.assert_close(weight.grad, gradient.to(dtype), rtol=0, atol=0) diff --git a/cosmos_framework/model/generator/algorithm/loss/fused_packed_cross_entropy_test.py b/cosmos_framework/model/generator/algorithm/loss/fused_packed_cross_entropy_test.py new file mode 100644 index 000000000..5eb38dedb --- /dev/null +++ b/cosmos_framework/model/generator/algorithm/loss/fused_packed_cross_entropy_test.py @@ -0,0 +1,60 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Packed fused loss must preserve independently evaluated conversation objectives.""" + +import pytest +import torch +import torch.nn.functional as F + +from cosmos_framework.model.generator.algorithm.loss import cross_entropy as ce + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + + +@pytest.mark.parametrize("exponent", [0.0, 0.5, 1.0]) +@pytest.mark.parametrize("recompute", [False, True]) +@pytest.mark.parametrize("all_ignored", [False, True]) +def test_packed_fused_objective_and_gradients_match_separate_conversations( + monkeypatch: pytest.MonkeyPatch, exponent: float, recompute: bool, all_ignored: bool +) -> None: + def scalar_ce(hidden: torch.Tensor, weight: torch.Tensor, labels: torch.Tensor, ignore_index: int) -> torch.Tensor: + # Only replace the CUDA kernel; exercise the real packing/objective/recompute path. + return F.cross_entropy(F.linear(hidden, weight), labels, ignore_index=ignore_index, reduction="sum") + + monkeypatch.setattr(ce, "_fused_linear_ce_sum", scalar_ce) + torch.manual_seed(715) + lengths = [4, 6, 3, 2] + hidden = torch.randn(1, sum(lengths), 5, requires_grad=True) + weight = torch.randn(13, 5, requires_grad=True) + labels = torch.tensor([[-100, 1, 2, 3, -100, 4, -100, 5, 6, 7, -100, -100, -100, -100, -100]]) + if all_ignored: + labels.fill_(-100) + boundaries = torch.tensor([0, 4, 10, 13, 15], dtype=torch.int32) + reference_hidden = hidden.detach().clone().requires_grad_() + reference_weight = weight.detach().clone().requires_grad_() + numerators: list[torch.Tensor] = [] + denominators: list[torch.Tensor] = [] + start = 0 + for length in lengths: + target = labels[0, start + 1 : start + length] + logits = F.linear(reference_hidden[0, start : start + length - 1], reference_weight) + count = (target != -100).sum().float() + numerators.append( + F.cross_entropy(logits, target, ignore_index=-100, reduction="sum") / count.clamp(min=1).pow(exponent) + ) + denominators.append(torch.where(count > 0, count.clamp(min=1).pow(1 - exponent), 0.0)) + start += length + denominator = torch.stack(denominators).sum() + expected = torch.stack(numerators).sum() / denominator.clamp(min=1) + actual, stats = ce.fused_weighted_cross_entropy_loss( + hidden, labels, weight, exponent=exponent, cu_seq_lens=boundaries, recompute=recompute + ) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(stats.objective_denominator, denominator) + torch.testing.assert_close( + ce.weighted_ce_denominator_from_labels(labels, exponent=exponent, cu_seq_lens=boundaries), denominator + ) + actual.backward() + expected.backward() + torch.testing.assert_close(hidden.grad, reference_hidden.grad) + torch.testing.assert_close(weight.grad, reference_weight.grad) diff --git a/cosmos_framework/model/generator/algorithm/loss/load_balancing_test.py b/cosmos_framework/model/generator/algorithm/loss/load_balancing_test.py index b63e7ef14..ac1ce7e0b 100644 --- a/cosmos_framework/model/generator/algorithm/loss/load_balancing_test.py +++ b/cosmos_framework/model/generator/algorithm/loss/load_balancing_test.py @@ -10,7 +10,8 @@ import torch import torch.distributed as dist import torch.multiprocessing as mp -from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard +from torch.distributed.tensor import DTensor from torch.distributed.tensor.device_mesh import DeviceMesh from torch.nn.parallel import DistributedDataParallel @@ -20,10 +21,11 @@ ) from cosmos_framework.model.generator.utils.load_balancing_stats import LBLMetadata, compute_sample_lbl_stats from cosmos_framework.data.generator.sequence_packing.runtime import prepare_sequence_pack_metadata - -pytestmark = [pytest.mark.L0, pytest.mark.CPU] +from cosmos_framework.utils.generator.parallelism import ParallelDims, fsdp_mesh +@pytest.mark.L0 +@pytest.mark.CPU def test_compute_sample_lbl_stats_preserves_packed_boundaries() -> None: routing_probabilities = torch.tensor( [ @@ -58,6 +60,8 @@ def test_compute_sample_lbl_stats_preserves_packed_boundaries() -> None: ) +@pytest.mark.L0 +@pytest.mark.CPU def test_compute_sample_lbl_stats_ignores_cp_padding_and_noncontiguous_samples() -> None: routing_probabilities = torch.tensor( [ @@ -83,6 +87,8 @@ def test_compute_sample_lbl_stats_ignores_cp_padding_and_noncontiguous_samples() torch.testing.assert_close(probability_sums, torch.tensor([[0.2, 0.8], [1.6, 0.4]])) # [B,E] +@pytest.mark.L0 +@pytest.mark.CPU def test_cp_shard_sample_stats_sum_to_unsharded_stats() -> None: routing_probabilities = torch.tensor( [ @@ -118,6 +124,8 @@ def test_cp_shard_sample_stats_sum_to_unsharded_stats() -> None: torch.testing.assert_close(rank0 + rank1, expected) +@pytest.mark.L0 +@pytest.mark.CPU def test_sequence_pack_sample_ids_follow_samples_across_multiple_splits() -> None: metadata = prepare_sequence_pack_metadata( sample_lens=[4, 3], @@ -131,6 +139,8 @@ def test_sequence_pack_sample_ids_follow_samples_across_multiple_splits() -> Non torch.testing.assert_close(metadata["_full_only_sample_ids"], torch.tensor([0, 1])) # [N_full] +@pytest.mark.L0 +@pytest.mark.CPU def test_sample_load_balancing_averages_packed_samples_equally() -> None: sample_router_prob_sum_per_expert = torch.tensor( [[[0.9, 0.1], [0.6, 2.4]]], @@ -165,7 +175,11 @@ def test_sample_load_balancing_averages_packed_samples_equally() -> None: def _distributed_metadata(routing_probabilities: torch.Tensor) -> LBLMetadata: """Build one-layer, one-sample router statistics from a local token shard.""" expert_indices = routing_probabilities.argmax(dim=-1, keepdim=True) # [N,1] - sample_ids = torch.zeros(routing_probabilities.shape[0], dtype=torch.long) # [N] + sample_ids = torch.zeros( + routing_probabilities.shape[0], + dtype=torch.long, + device=routing_probabilities.device, + ) # [N] sample_counts, sample_tokens, sample_probability_sums = compute_sample_lbl_stats( routing_probabilities, expert_indices, @@ -176,7 +190,7 @@ def _distributed_metadata(routing_probabilities: torch.Tensor) -> LBLMetadata: num_tokens_per_expert=sample_counts.sum(dim=0, keepdim=True), # [1,E] num_tokens=sample_tokens.sum(dim=0, keepdim=True), # [1,1] mean_router_prob_per_expert=routing_probabilities.mean(dim=0, keepdim=True), # [1,E] - top_k=torch.ones((1, 1), dtype=torch.long), # [1,1] + top_k=torch.ones((1, 1), dtype=torch.long, device=routing_probabilities.device), # [1,1] sample_num_tokens_per_expert=sample_counts.unsqueeze(0), # [1,1,E] sample_num_tokens=sample_tokens.unsqueeze(0), # [1,1,1] sample_router_prob_sum_per_expert=sample_probability_sums.unsqueeze(0), # [1,1,E] @@ -202,8 +216,7 @@ def forward(self, features: torch.Tensor) -> LBLMetadata: return _distributed_metadata(torch.softmax(features @ self.weight, dim=-1)) # [N,E] -def _distributed_cp_gradient_worker(rank: int, rendezvous: str) -> None: - world_size = 2 +def _distributed_cp_gradient_worker(rank: int, rendezvous: str, world_size: int, cp_size: int) -> None: torch.set_num_threads(1) dist.init_process_group( "gloo", @@ -213,12 +226,20 @@ def _distributed_cp_gradient_worker(rank: int, rendezvous: str) -> None: timeout=timedelta(seconds=90), ) try: - mesh = init_device_mesh("cpu", (world_size,), mesh_dim_names=("cp",)) - logical_features = torch.tensor( - [[1.0, 0.2], [0.5, -0.7], [-0.2, 0.8], [0.9, -0.4]], - dtype=torch.float64, - ) - local_features = logical_features.chunk(world_size)[rank] # [N/CP,D] + parallel_dims = ParallelDims(world_size=world_size, dp_shard=world_size, cp=cp_size) + parallel_dims.build_meshes("cpu") + dp_mesh = parallel_dims.dp_mesh + cp_mesh = parallel_dims.cp_mesh + assert dp_mesh is not None and dp_mesh.size() == world_size + assert cp_mesh is not None and cp_mesh.size() == cp_size + sample_count = world_size // cp_size + logical_features = [ + torch.tensor([[1.0, 0.2], [0.5, -0.7], [-0.2, 0.8], [0.9, -0.4]], dtype=torch.float64), + torch.tensor([[-0.6, 0.3], [0.2, 1.1], [0.7, 0.4], [-0.8, -0.5]], dtype=torch.float64), + ][:sample_count] + sample_rank = rank // cp_size + cp_rank = rank % cp_size + local_features = logical_features[sample_rank].chunk(cp_size)[cp_rank] # [N/CP,D] router = _DistributedRouter() reference = copy.deepcopy(router) @@ -230,42 +251,152 @@ def _distributed_cp_gradient_worker(rank: int, rendezvous: str) -> None: distributed_router(local_features), coeff=1.0, method="sample", - device_mesh=mesh, - context_parallel_mesh=mesh, + device_mesh=dp_mesh, + context_parallel_mesh=cp_mesh, correct_cp_gradients=correct_cp_gradients, ) assert local_loss is not None local_loss.backward() reference.zero_grad(set_to_none=True) - expected_loss = compute_load_balancing_loss( - reference(logical_features), - coeff=1.0, - method="sample", - device_mesh=None, - ) - assert expected_loss is not None + reference_losses: list[torch.Tensor] = [] + for features in logical_features: + reference_loss = compute_load_balancing_loss( + reference(features), + coeff=1.0, + method="sample", + device_mesh=None, + ) + assert reference_loss is not None + reference_losses.append(reference_loss) + expected_loss = torch.stack(reference_losses).mean() expected_loss.backward() - torch.testing.assert_close(local_loss, expected_loss) + torch.testing.assert_close(local_loss, reference_losses[sample_rank]) expected_gradient = reference.weight.grad assert router.weight.grad is not None and expected_gradient is not None if not correct_cp_gradients: - expected_gradient = expected_gradient / world_size + expected_gradient = expected_gradient / cp_size torch.testing.assert_close(router.weight.grad, expected_gradient) finally: dist.destroy_process_group() +@pytest.mark.L0 +@pytest.mark.CPU def test_sample_load_balancing_corrects_cp_gradients_without_changing_loss(tmp_path: Path) -> None: mp.spawn( _distributed_cp_gradient_worker, - args=((tmp_path / "rendezvous").as_uri(),), + args=((tmp_path / "two_rank_rendezvous").as_uri(), 2, 2), nprocs=2, join=True, ) +@pytest.mark.L0 +@pytest.mark.CPU +def test_sample_load_balancing_cp_correction_with_separate_dp_and_cp_meshes(tmp_path: Path) -> None: + mp.spawn( + _distributed_cp_gradient_worker, + args=((tmp_path / "separate_mesh_rendezvous").as_uri(), 4, 2), + nprocs=4, + join=True, + ) + + +def _fsdp_cp_gradient_worker(rank: int, rendezvous: str) -> None: + world_size = 4 + cp_size = 2 + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", + init_method=rendezvous, + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=90), + ) + try: + parallel_dims = ParallelDims(world_size=world_size, dp_shard=world_size, cp=cp_size) + parallel_dims.build_meshes("cuda") + dp_mesh = parallel_dims.dp_mesh + cp_mesh = parallel_dims.cp_mesh + shard_mesh = fsdp_mesh(parallel_dims) + assert dp_mesh is not None and dp_mesh.size() == world_size + assert cp_mesh is not None and cp_mesh.size() == cp_size + assert shard_mesh is not None and shard_mesh.size() == world_size + + device = torch.device("cuda", rank) + logical_features = [ + torch.tensor([[1.0, 0.2], [0.5, -0.7], [-0.2, 0.8], [0.9, -0.4]], dtype=torch.float64, device=device), + torch.tensor( + [[-0.6, 0.3], [0.2, 1.1], [0.7, 0.4], [-0.8, -0.5]], + dtype=torch.float64, + device=device, + ), + ] + sample_rank = rank // cp_size + cp_rank = rank % cp_size + local_features = logical_features[sample_rank].chunk(cp_size)[cp_rank] # [N/CP,D] + + router = _DistributedRouter().to(device) + reference = copy.deepcopy(router) + fully_shard(router, mesh=shard_mesh) + + for correct_cp_gradients in (False, True): + router.zero_grad(set_to_none=True) + local_loss = compute_load_balancing_loss( + router(local_features), + coeff=1.0, + method="sample", + device_mesh=dp_mesh, + context_parallel_mesh=cp_mesh, + correct_cp_gradients=correct_cp_gradients, + ) + assert local_loss is not None + local_loss.backward() + + reference.zero_grad(set_to_none=True) + reference_losses: list[torch.Tensor] = [] + for features in logical_features: + reference_loss = compute_load_balancing_loss( + reference(features), + coeff=1.0, + method="sample", + device_mesh=None, + ) + assert reference_loss is not None + reference_losses.append(reference_loss) + torch.stack(reference_losses).mean().backward() + + torch.testing.assert_close(local_loss, reference_losses[sample_rank]) + distributed_gradient = router.weight.grad + expected_gradient = reference.weight.grad + assert isinstance(distributed_gradient, DTensor) and expected_gradient is not None + full_distributed_gradient = distributed_gradient.full_tensor() + if not correct_cp_gradients: + expected_gradient = expected_gradient / cp_size + torch.testing.assert_close(full_distributed_gradient, expected_gradient) + finally: + dist.destroy_process_group() + + +@pytest.mark.L1 +@pytest.mark.GPU +def test_sample_load_balancing_cp_correction_matches_reference_after_fsdp_reduce_scatter(tmp_path: Path) -> None: + required_gpus = 4 + if torch.cuda.device_count() < required_gpus: + pytest.skip(f"requires {required_gpus} visible GPUs") + + mp.spawn( + _fsdp_cp_gradient_worker, + args=((tmp_path / "fsdp_rendezvous").as_uri(),), + nprocs=required_gpus, + join=True, + ) + + +@pytest.mark.L0 +@pytest.mark.CPU def test_sample_load_balancing_scales_rank_means_for_fsdp_average( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -321,6 +452,8 @@ def _fake_sum_over_mesh(local_tensor: torch.Tensor, mesh: DeviceMesh) -> torch.T torch.testing.assert_close(rank1_local_mean.grad, torch.tensor([0.75])) # [num_layers] +@pytest.mark.L0 +@pytest.mark.CPU def test_sample_load_balancing_requires_sample_metadata() -> None: metadata = LBLMetadata( num_tokens_per_expert=torch.tensor([[1, 1]]), # [num_layers,num_experts] @@ -338,6 +471,8 @@ def test_sample_load_balancing_requires_sample_metadata() -> None: ) +@pytest.mark.L0 +@pytest.mark.CPU def test_sample_load_balancing_returns_zero_without_valid_samples() -> None: sample_router_prob_sum_per_expert = torch.zeros((1, 2, 2), requires_grad=True) # [num_layers,num_samples,E] metadata = LBLMetadata( @@ -365,6 +500,8 @@ def test_sample_load_balancing_returns_zero_without_valid_samples() -> None: ) +@pytest.mark.L0 +@pytest.mark.CPU def test_sample_load_balancing_matches_local_global_top_k_scaling() -> None: top_k = torch.tensor([[2]]) # [num_layers,1] metadata = LBLMetadata( diff --git a/cosmos_framework/model/generator/context_parallel.py b/cosmos_framework/model/generator/context_parallel.py new file mode 100644 index 000000000..7b21a5ce5 --- /dev/null +++ b/cosmos_framework/model/generator/context_parallel.py @@ -0,0 +1,228 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Shared Ulysses communication and sequence ownership for native VLMs.""" + +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist +from torch import nn + + +@dataclass(frozen=True) +class ContextParallelBatchPolicy: + """Preserve the native position layout and input-sharding stage. + + ``position_axes=None`` denotes [B,T]; an integer denotes [A,B,T]. + Packed positions are prepared by the model's native position routine. + ``keep_full_inputs`` also retains text-only tokens until decoder entry. + Otherwise only enabled multimodal batches defer sharding until fusion. + """ + + position_axes: int | None = None + allow_multimodal: bool = False + keep_full_inputs: bool = False + requires_packed: bool = False + + +@dataclass(frozen=True) +class ContextParallelModelSpec: + """Explicit native-model bindings for ``parallelize_vlm.apply_cp``. + + Module references are metadata, not additional registered parameters. + Integrations validate native variants before returning this spec. Merely + exposing q/k/v projections does not establish CP support. Position + generation, attention kernels and loss precision stay with the model. + """ + + decoder: nn.Module + attention_modules: tuple[nn.Module, ...] + text_config: Any + fused_ce_enabled: bool + position_axes: int | None + keep_full_inputs: bool = False + requires_packed: bool = False + require_default_rope: bool = False + decoder_packing_key: str | None = None + fusion_module: nn.Module | None = None + + +def _all_to_all(tensor: torch.Tensor, group: dist.ProcessGroup, scatter_dim: int, gather_dim: int) -> torch.Tensor: + size = dist.get_world_size(group) + if tensor.shape[scatter_dim] % size: + raise ValueError("Ulysses scatter dimension must be divisible by CP size") + chunks = tensor.chunk(size, dim=scatter_dim) + shape = chunks[0].shape + packed = torch.cat([chunk.contiguous().view(-1) for chunk in chunks]) + received = torch.empty_like(packed) + dist.all_to_all_single(received, packed, group=group) + return torch.cat([chunk.view(shape) for chunk in received.chunk(size)], dim=gather_dim).contiguous() + + +class _SequenceAllToAll(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, tensor: torch.Tensor, group: dist.ProcessGroup, scatter_dim: int, gather_dim: int + ) -> torch.Tensor: + ctx.group = group + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + return _all_to_all(tensor, group, scatter_dim, gather_dim) + + @staticmethod + def backward(ctx: Any, gradient: torch.Tensor) -> tuple[torch.Tensor, None, None, None]: + return _all_to_all(gradient, ctx.group, ctx.gather_dim, ctx.scatter_dim), None, None, None + + +def sequence_all_to_all( + tensor: torch.Tensor, group: dist.ProcessGroup, *, scatter_dim: int, gather_dim: int +) -> torch.Tensor: + """Swap sequence shards for head shards; the backward performs the inverse.""" + return _SequenceAllToAll.apply(tensor, group, scatter_dim, gather_dim) + + +def shard_text_batch( + data: dict[str, Any], group: dist.ProcessGroup, *, position_axes: int | None = None +) -> dict[str, Any]: + """Slice replicated text inputs, global positions, and pre-shifted next-token labels. + + The loader must replicate each logical sample across adjacent CP ranks. Labels + are shifted BEFORE slicing, preserving the target at every shard boundary. + Every CP shard keeps the full model; FSDP still averages gradients over WORLD. + """ + size = dist.get_world_size(group) + if data.get("context_parallel_size") != size: + raise ValueError("CP requires a loader that replicates each sample within its CP group") + packed = data.get("cu_seq_lens_q") is not None + if data.get("true_packing", False): + raise ValueError("Prepare global packed positions and boundaries before CP slicing") + if any( + data.get(key) is not None + for key in ( + "pixel_values", + "pixel_values_videos", + "image_grid_thw", + "video_grid_thw", + "inputs_embeds", + "past_key_values", + "cache_position", + "audio_values", + "input_features", + ) + ): + raise NotImplementedError("Text CP requires token inputs without media or KV cache") + if not packed and data.get("position_ids") is not None: + raise ValueError("Unpacked text CP computes its own global positions") + if packed and data.get("cu_seq_lens_q") is not data.get("cu_seq_lens_k"): + raise ValueError("Packed CP requires shared global Q/KV boundaries") + ids, labels, mask = data["input_ids"], data["labels"], data["attention_mask"] + if ids.ndim != 2 or ids.shape[0] != 1 or labels.shape != ids.shape or mask.shape != ids.shape: + raise ValueError("Text CP requires aligned [1,T] input_ids, labels and attention_mask") + if ids.shape[1] % size: + raise ValueError("Global sequence length must be divisible by CP size") + if bool(((mask != 0) & (mask != 1)).any()) or bool((mask[:, 1:] > mask[:, :-1]).any()): + raise ValueError("Text CP accepts right padding only") + if bool((labels[mask == 0] != -100).any()): + raise ValueError("Padding labels must be ignored") + local_length = ids.shape[1] // size + start = dist.get_rank(group) * local_length + section = slice(start, start + local_length) + targets = torch.cat([labels[:, 1:], torch.full_like(labels[:, :1], -100)], dim=1) + positions = data.get("position_ids") if packed else (mask.long().cumsum(-1) - 1).clamp(min=0) + if not packed and position_axes is not None: + positions = positions.unsqueeze(0).expand(position_axes, -1, -1) + expected_shape = (1, ids.shape[1]) if position_axes is None else (position_axes, 1, ids.shape[1]) + if not isinstance(positions, torch.Tensor) or positions.shape != expected_shape: + raise ValueError("Packed CP requires complete native per-conversation positions") + output = dict(data) + output.update( + input_ids=ids[:, section].contiguous(), + labels=targets[:, section].contiguous(), + attention_mask=mask[:, section].contiguous(), + position_ids=positions[..., section].contiguous(), + ) + return output + + +def context_parallel_router_mean( + local_mean: torch.Tensor, local_tokens: torch.Tensor, group: dist.ProcessGroup +) -> torch.Tensor: + """Scale a shard's router mean so WORLD averaging recovers each full sample mean.""" + sequence_tokens = local_tokens.detach().clone() + dist.all_reduce(sequence_tokens, group=group) + return local_mean * (local_tokens * dist.get_world_size(group) / sequence_tokens.clamp_min(1)) + + +def prepare_context_parallel_batch( + data: dict[str, Any], group: dist.ProcessGroup, *, policy: ContextParallelBatchPolicy +) -> dict[str, Any]: + """Keep media and full tokens until fusion, but shift/shard labels before CE. + + The caller must supply identical preprocessed media within each CP group. + This first version does not distribute raw-media decoding or vision encoding. + """ + if policy.requires_packed and data.get("cu_seq_lens_q") is None: + raise ValueError("This native CP integration requires true-packed inputs") + media_keys = ("pixel_values", "pixel_values_videos", "image_grid_thw", "video_grid_thw") + has_media = any(data.get(key) is not None for key in media_keys) + if not policy.keep_full_inputs and (not has_media or not policy.allow_multimodal): + return shard_text_batch(data, group, position_axes=policy.position_axes) + if data.get("cache_position") is not None: + raise NotImplementedError("Multimodal CP requires a complete training sequence without cache") + # Reuse shape, padding, replication, unsupported-input and boundary-target checks. + local = shard_text_batch( + {key: value for key, value in data.items() if key not in media_keys}, group, position_axes=policy.position_axes + ) + output = dict(data) + output["labels"] = local["labels"] + output["context_parallel_token_mask"] = local["attention_mask"] + return output + + +def shard_decoder_inputs( + inputs: dict[str, Any], group: dist.ProcessGroup, *, position_axes: int | None, clone_embeddings: bool = False +) -> dict[str, Any]: + """Shard after complete media fusion and native global position calculation. + + Visual rows in DeepStack tensors are compressed (text rows are absent), so + select by global visual-token positions rather than the decoder row slice. + Empty selections retain their autograd connection to every vision module. + FSDP WORLD averaging and the CP-scaled fused objective perform gradient + aggregation; an additional vision-gradient all-reduce would double count. + + ``clone_embeddings`` preserves integrations whose decoder modifies the + fused embedding buffer in place. Other integrations retain contiguous views. + Global packed boundaries and all nonsequence arguments pass through intact. + """ + size = dist.get_world_size(group) + hidden, positions = inputs.get("inputs_embeds"), inputs.get("position_ids") + if not isinstance(hidden, torch.Tensor) or hidden.ndim != 3 or hidden.shape[0] != 1: + raise ValueError("CP requires complete fused embeddings shaped [1,T,H]") + length = hidden.shape[1] + expected_shape = (1, length) if position_axes is None else (position_axes, 1, length) + if length % size or not isinstance(positions, torch.Tensor) or positions.shape != expected_shape: + raise ValueError("CP requires divisible tokens and complete native positions") + local_length = length // size + start = dist.get_rank(group) * local_length + section = slice(start, start + local_length) + output = dict(inputs) + output["inputs_embeds"] = hidden[:, section].clone() if clone_embeddings else hidden[:, section].contiguous() + output["position_ids"] = positions[..., section].contiguous() + for key in ("attention_mask", "padding_mask", "visual_pos_masks"): + value = output.get(key) + if value is not None: + if not isinstance(value, torch.Tensor) or value.shape != (1, length): + raise ValueError(f"CP requires a full [1,T] {key}") + output[key] = value[:, section].contiguous() + features = inputs.get("deepstack_visual_embeds") + if features is not None: + visual = inputs.get("visual_pos_masks") + if not isinstance(visual, torch.Tensor): + raise ValueError("Compressed visual features require the full visual token mask") + visual_positions = visual[0].nonzero(as_tuple=True)[0] + selected = (visual_positions >= start) & (visual_positions < start + local_length) + if any(value.shape[0] != visual_positions.numel() for value in features): + raise ValueError("DeepStack rows must match all global visual tokens") + output["deepstack_visual_embeds"] = [value[selected].contiguous() for value in features] + return output diff --git a/cosmos_framework/model/generator/hf_model.py b/cosmos_framework/model/generator/hf_model.py index ac4889dba..cdc20ade1 100644 --- a/cosmos_framework/model/generator/hf_model.py +++ b/cosmos_framework/model/generator/hf_model.py @@ -34,13 +34,15 @@ from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoModelForImageTextToText, AutoTokenizer import cosmos_framework.model.generator.reasoner.cosmos3_edge # noqa: F401 registers cosmos3_edge with transformers Auto classes +from cosmos_framework.utils import log +from cosmos_framework.model.generator.context_parallel import ContextParallelBatchPolicy, ContextParallelModelSpec +from cosmos_framework.model.generator.reasoner.qwen3_vl_fused_ce import configure_qwen3_vl_fused_ce from cosmos_framework.model.generator.reasoner.qwen35_caption import ( Qwen35CaptionLoss, configure_qwen35_caption_model, named_parameters_with_qwen35_decay, ) from cosmos_framework.model.generator.utils.safetensors_loader import load_language_model, load_vlm_model -from cosmos_framework.utils import log from cosmos_framework.utils.generator.input_probe import maybe_dump_pre_forward from cosmos_framework.utils.generator.parallelism import ParallelDims @@ -204,6 +206,21 @@ def _tensor_names_to_skip_for(model_type: str) -> list[str]: return table.get(model_type, []) +def configure_sound_und_encoder_trainability(sound_und_model: nn.Module, *, freeze_encoder: bool) -> None: + """Apply the coarse audio policy plus backend-specific unused-parameter freezes.""" + encoder = getattr(sound_und_model, "encoder", None) + if not isinstance(encoder, nn.Module): + raise TypeError("sound_und_model.encoder must be an nn.Module") + encoder.requires_grad_(not freeze_encoder) + configure = getattr(sound_und_model, "configure_encoder_trainability", None) + if configure is not None: + if not callable(configure): + raise TypeError("configure_encoder_trainability must be callable") + configure(train_encoder=not freeze_encoder) + if freeze_encoder: + encoder.eval() + + class HFModel(nn.Module): """Minimal HF model wrapper for the vfm/ unified VLM training path. @@ -232,6 +249,28 @@ class (e.g. ``Qwen3VLForConditionalGeneration``), which exposes ``.logits`` ``vfm/models/parallelize_vlm.py``. """ + # Validation video cache state; __init__ sets these, the defaults cover wrappers built + # with __new__ (tests) so forward/train do not depend on __init__ having run. + _cosmos_validation_video_cache_keys = None + _cosmos_validation_video_cache_active = False + _cosmos_validation_video_feature_cache = None + + def context_parallel_spec(self) -> ContextParallelModelSpec: + """Bind validated native modules; shared CP owns communication and slicing.""" + if self.hf_config.model_type not in {"qwen3_vl", "qwen3_vl_moe"}: + raise NotImplementedError("This native model has no validated context-parallel integration") + # Only the validated native dense/MoE model interface is patched. + decoder = self.model.model.language_model + return ContextParallelModelSpec( + decoder=decoder, + attention_modules=tuple(layer.self_attn for layer in decoder.layers), + text_config=self.model.config.text_config, + fused_ce_enabled=bool(getattr(self.model, "_qwen3_vl_fused_ce", False)), + position_axes=3, + require_default_rope=True, + fusion_module=self.model.model, + ) + def named_parameters( self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True ) -> Iterator[tuple[str, nn.Parameter]]: @@ -248,9 +287,15 @@ def __init__( configured_model_name_or_path: str | None = None, qwen35_fp32_recurrent_a_log: bool = False, enable_fused_weighted_ce: bool = False, + recompute_fused_weighted_ce: bool = False, weighted_ce_exponent: float = 0.5, + qwen3_vl_moe_fixed_sum: bool = False, + qwen3_vl_vision_attention_max_padded_tokens: int = 0, + qwen3_vl_decoder_attention_max_padded_tokens: int = 0, + qwen3_vl_accumulate_head_grad_fp32: bool = True, ): super().__init__() + self.context_parallel_batch_policy: ContextParallelBatchPolicy | None = None self.model_name_or_path = model_name_or_path self.configured_model_name_or_path = configured_model_name_or_path or model_name_or_path self.sound_und = sound_und @@ -262,8 +307,31 @@ def __init__( raise ValueError("sound_und_config is required when sound_und=True") hf_config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=trust_remote_code) self.hf_config = hf_config - if (qwen35_fp32_recurrent_a_log or enable_fused_weighted_ce) and hf_config.model_type != "qwen3_5": - raise ValueError("Caption FP32/fused loss options require dense qwen3_5") + if sound_und and enable_fused_weighted_ce and hf_config.model_type in {"qwen3_vl", "qwen3_vl_moe"}: + raise ValueError("Qwen3-VL fused weighted CE does not support sound_und=True") + if qwen3_vl_decoder_attention_max_padded_tokens < 0: + raise ValueError("Decoder attention padded-token budget must be nonnegative") + if qwen3_vl_decoder_attention_max_padded_tokens and ( + hf_config.model_type not in {"qwen3_vl", "qwen3_vl_moe"} or attn_implementation != "cosmos" + ): + raise ValueError("Bounded decoder varlen attention requires Qwen3-VL with the cosmos backend") + if qwen3_vl_vision_attention_max_padded_tokens < 0: + raise ValueError("Vision attention padded-token budget must be nonnegative") + if qwen3_vl_vision_attention_max_padded_tokens and ( + hf_config.model_type not in {"qwen3_vl", "qwen3_vl_moe"} or attn_implementation != "cosmos" + ): + raise ValueError("Bounded vision varlen attention requires Qwen3-VL with the cosmos backend") + if qwen3_vl_moe_fixed_sum and hf_config.model_type != "qwen3_vl_moe": + raise ValueError("Fixed MoE summation requires qwen3_vl_moe") + if qwen35_fp32_recurrent_a_log and hf_config.model_type != "qwen3_5": + raise ValueError("Caption FP32 recurrent decay requires dense qwen3_5") + if enable_fused_weighted_ce and hf_config.model_type not in {"qwen3_5", "qwen3_vl", "qwen3_vl_moe"}: + raise ValueError("Fused weighted CE requires qwen3_5, qwen3_vl, or qwen3_vl_moe") + + if recompute_fused_weighted_ce and ( + not enable_fused_weighted_ce or hf_config.model_type not in {"qwen3_vl", "qwen3_vl_moe"} + ): + raise ValueError("CE recomputation requires Qwen3-VL with enable_fused_weighted_ce=True") # Register cosmos before from_config validates it. Gated so non-cosmos # paths don't import cosmos_framework.model.attention. @@ -305,9 +373,9 @@ def __init__( ) if sound_und: - from cosmos_framework.data.generator.processors.audio_utils import add_reasoner_audio_special_tokens from cosmos_framework.model.generator.reasoner.audio.registry import get_audio_encoder_backend from cosmos_framework.model.generator.reasoner.audio.utils import patch_reasoner_audio_forward + from cosmos_framework.data.generator.processors.audio_utils import add_reasoner_audio_special_tokens audio_backend = get_audio_encoder_backend(sound_und_config.audio_encoder_type) @@ -341,8 +409,9 @@ def __init__( sound_und_config.projection_hidden_size, reasoner_hidden_size, ) - self.model.sound_und_model.encoder.requires_grad_(False) - self.model.sound_und_model.encoder.eval() + configure_sound_und_encoder_trainability( + self.model.sound_und_model, freeze_encoder=sound_und_config.freeze_encoder + ) if sound_und_config.freeze_projector: self.model.sound_und_model.projector.requires_grad_(False) self.model.sound_und_model.projector.eval() @@ -402,6 +471,13 @@ def __init__( patch_qwen3_vl_forward(self.model.model) log.info(f"HFModel: applied patch_qwen3_vl_forward ({hf_config.model_type}) for text-only batch support") + if enable_fused_weighted_ce: + configure_qwen3_vl_fused_ce( + self.model, + exponent=weighted_ce_exponent, + recompute=recompute_fused_weighted_ce, + accumulate_head_grad_fp32=qwen3_vl_accumulate_head_grad_fp32, + ) # Swap HF's per-expert Python loop for the fused grouped_mm expert kernel. HF's loop # is never the faster choice on the GPUs this trains on, and the patch reuses the @@ -412,6 +488,13 @@ def __init__( n_moe_blocks = patch_qwen3_vl_moe_grouped_mm_experts(self.model) log.info(f"HFModel: applied grouped_mm experts to {n_moe_blocks} MoE block(s)") + if qwen3_vl_moe_fixed_sum: + from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.fixed_sum import ( + configure_qwen3_vl_moe_fixed_sum, + ) + + n_fixed = configure_qwen3_vl_moe_fixed_sum(self.model) + log.info(f"HFModel: enabled fixed-order FP32 summation in {n_fixed} MoE block(s)") # Give the vision tower a single varlen attention call per block. HF only does that for # flash_attention_2 and otherwise splits the packed patches per image, which costs a @@ -421,7 +504,9 @@ def __init__( if hf_config.model_type in ("qwen3_vl", "qwen3_vl_moe") and attn_implementation == "cosmos": from cosmos_framework.utils.generator.monkey_patch import patch_qwen3_vl_vision_varlen_attention - n_vision_attns = patch_qwen3_vl_vision_varlen_attention(self.model) + n_vision_attns = patch_qwen3_vl_vision_varlen_attention( + self.model, max_padded_tokens=qwen3_vl_vision_attention_max_padded_tokens + ) log.info(f"HFModel: applied varlen attention to {n_vision_attns} vision attention module(s)") elif hf_config.model_type == "qwen3_vl" and hasattr(self.model, "model"): @@ -497,7 +582,8 @@ def train(self, mode: bool = True) -> "HFModel": self.clear_validation_video_feature_cache() super().train(mode) if self.sound_und: - self.model.sound_und_model.encoder.eval() + if self.sound_und_config.freeze_encoder: + self.model.sound_und_model.encoder.eval() if self.sound_und_config.freeze_projector: self.model.sound_und_model.projector.eval() return self @@ -657,6 +743,8 @@ def load_weights( "cu_seq_lens_k", "max_length_q", "max_length_k", + # CPU-planned groups of complete packed conversations for bounded workspaces. + "decoder_attention_chunks", } ) diff --git a/cosmos_framework/model/generator/mot/activation_marks_test.py b/cosmos_framework/model/generator/mot/activation_marks_test.py index 299386e97..d74047c27 100644 --- a/cosmos_framework/model/generator/mot/activation_marks_test.py +++ b/cosmos_framework/model/generator/mot/activation_marks_test.py @@ -23,6 +23,7 @@ from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts from cosmos_framework.model.attention import attention +from cosmos_framework.utils.helper_test import RunIf from cosmos_framework.model.generator.mot.activation_marks import ( MARK_OP_QUALNAME, enable_marking, @@ -32,7 +33,6 @@ reset_marking_for_tests, ) from cosmos_framework.model.generator.mot.parallelize_unified_mot import make_selective_ac_policy -from cosmos_framework.utils.helper_test import RunIf HEADS = 8 KV_HEADS = 2 @@ -265,10 +265,10 @@ def test_the_attention_op_regex_covers_every_backend_that_can_be_selected() -> N kernel fails here instead of in a training run's memory profile. cuDNN is out of scope by design: it rejects varlen and every marked call site is varlen. """ + from cosmos_framework.model.attention.backends import get_backend_list from cosmos_framework.configs.base.defaults.activation_checkpointing import ( ATTENTION_FORWARD_OPS_REGEX, ) - from cosmos_framework.model.attention.backends import get_backend_list patterns = [re.compile(pattern) for pattern in ATTENTION_FORWARD_OPS_REGEX] diff --git a/cosmos_framework/model/generator/mot/causal_attention.py b/cosmos_framework/model/generator/mot/causal_attention.py index 15d557dce..927998dd3 100644 --- a/cosmos_framework/model/generator/mot/causal_attention.py +++ b/cosmos_framework/model/generator/mot/causal_attention.py @@ -14,23 +14,13 @@ import torch from torch.nn.attention.flex_attention import BlockMask -from cosmos_framework.configs.base.defaults.replay_attention import TeacherForcingReplayPolicyConfig -from cosmos_framework.data.generator.sequence_packing.runtime import ( - SequencePack, - drop_pad_segment, - from_mode_splits, - from_und_gen_splits, - get_caption_seq_offsets, - get_causal_seq, - get_full_only_seq, - get_gen_seq, -) from cosmos_framework.model.attention import ( attention, multi_dimensional_attention, multi_dimensional_attention_varlen, ) from cosmos_framework.model.attention.masks import CausalType +from cosmos_framework.configs.base.defaults.replay_attention import TeacherForcingReplayPolicyConfig from cosmos_framework.model.generator.mot.attention import SplitInfo from cosmos_framework.model.generator.mot.attention import dispatch_attention as vfm_dispatch_attention from cosmos_framework.model.generator.mot.flex_attention import FlexBackend, flex_attention @@ -55,11 +45,21 @@ write_current_kv_to_static_buffer, zero_null_action_values, ) -from cosmos_framework.model.generator.utils.memory import KVToStore, MemoryValue +from cosmos_framework.model.generator.utils.memory import ARTextLayerCache, KVToStore, MemoryValue from cosmos_framework.model.generator.utils.rolling_kv.rolling_prompt import ( RollingPromptLayout, RollingTextSinkLayer, ) +from cosmos_framework.data.generator.sequence_packing.runtime import ( + SequencePack, + drop_pad_segment, + from_mode_splits, + from_und_gen_splits, + get_caption_seq_offsets, + get_causal_seq, + get_full_only_seq, + get_gen_seq, +) class ConcatenateAttentionsBridge(torch.autograd.Function): @@ -416,6 +416,7 @@ def multiview_attention_with_memory( replay_maskless_plan: ReplayMasklessPlan | None = None, text_sink_layer: RollingTextSinkLayer | None = None, prompt_layout: RollingPromptLayout | None = None, + text_cache: ARTextLayerCache | None = None, ) -> SequencePack: """Run Flex or maskless multiview replay with an optional key-only K/V suffix.""" if (memory_k is None) != (memory_v is None): @@ -449,20 +450,25 @@ def multiview_attention_with_memory( torch._check(max_causal_len <= causal_k.shape[0]) use_dont_care_mask = causal_q_offsets is causal_k_offsets - causal_res = attention( - causal_q.unsqueeze(0), # [1,N_und,H,D] - causal_k.unsqueeze(0), # [1,N_und,H,D] - causal_v.unsqueeze(0), # [1,N_und,H,D] - cumulative_seqlen_Q=causal_q_offsets, - cumulative_seqlen_KV=causal_k_offsets, - max_seqlen_Q=max_causal_len, - max_seqlen_KV=max_causal_len, - is_causal=True, - causal_type=CausalType.DontCare if use_dont_care_mask else CausalType.TopLeft, - ) # [1,N_und,H,D] - if not isinstance(causal_res, torch.Tensor): - raise TypeError("Two-way causal attention must return a tensor.") - causal_out = causal_res.squeeze(0).flatten(-2, -1) # [N_und,H*D] + if text_cache is not None and text_cache.attention_output is not None: + causal_out = text_cache.attention_output # [N_und,H*D] + else: + causal_res = attention( + causal_q.unsqueeze(0), # [1,N_und,H,D] + causal_k.unsqueeze(0), # [1,N_und,H,D] + causal_v.unsqueeze(0), # [1,N_und,H,D] + cumulative_seqlen_Q=causal_q_offsets, + cumulative_seqlen_KV=causal_k_offsets, + max_seqlen_Q=max_causal_len, + max_seqlen_KV=max_causal_len, + is_causal=True, + causal_type=CausalType.DontCare if use_dont_care_mask else CausalType.TopLeft, + ) # [1,N_und,H,D] + if not isinstance(causal_res, torch.Tensor): + raise TypeError("Two-way causal attention must return a tensor.") + causal_out = causal_res.squeeze(0).flatten(-2, -1) # [N_und,H*D] + if text_cache is not None: + text_cache.attention_output = causal_out und_v = causal_v # [N_und,H,D] gen_k, _ = get_full_only_seq(packed_key_normalized) # [N_gen,H,D], [B+1] @@ -1761,26 +1767,34 @@ def three_way_attention_with_memory( video_out = torch.nn.functional.pad(video_out, (0, 0, 0, pad_len)) # --- text self-attention --- - # Varlen self-attention over the live text in the current pack. The + # Varlen self-attention over the live text in the current pack. The # text Q/K/V tensors are padded to a constant length (compile-stable - # shape); the varlen offsets restrict the kernel to the real - # (non-padded) range. Output positions beyond the real range are - # discarded by ``from_mode_splits`` so they never reach the loss. + # shape). Keep K/V restricted to the real range, but evaluate every + # allocated query row: Hopper NATTEN leaves padded-query backward + # workspace uninitialized when Q offsets cover only the real prefix, + # which can introduce NaNs into real K gradients under torch.compile. + # Explicitly zero the padded output rows before returning the pack. # - # Reuse ``text_kv_offsets``, which should be the same as text Q offsets, and - # is already pre-clamped to >=1 by ``KVCacheTrainMemoryState.init()`` - # to avoid kernel failure in case length == 0. + # ``KVCacheTrainMemoryState.init()`` normally clamps empty K/V to >=1. + # Do that locally even when the global clamp is disabled: the added + # queries need a finite dummy attention result before we mask them out. + # Top-left causality ensures that adding padded queries cannot change + # the real caption's outputs or gradients. text_q, _ = get_causal_seq(packed_query_states) text_k, _ = get_causal_seq(packed_key_states) text_v, _ = get_causal_seq(packed_value_states) padded_text_len = text_q.shape[0] + text_q_offsets = text_kv_offsets.new_tensor([0, padded_text_len]) # [2] + text_sa_kv_offsets = torch.stack( # [2] + (text_kv_offsets[0], text_kv_offsets[1].clamp_min(min(padded_text_len, 1))) + ) text_res = attention( text_q.unsqueeze(0), text_k.unsqueeze(0), text_v.unsqueeze(0), - cumulative_seqlen_Q=text_kv_offsets, - cumulative_seqlen_KV=text_kv_offsets, + cumulative_seqlen_Q=text_q_offsets, + cumulative_seqlen_KV=text_sa_kv_offsets, max_seqlen_Q=padded_text_len, max_seqlen_KV=padded_text_len, is_causal=True, @@ -1788,6 +1802,10 @@ def three_way_attention_with_memory( backend="natten", # Sanity check; this should be the only option. ) assert isinstance(text_res, torch.Tensor) + real_text_rows = ( + torch.arange(padded_text_len, device=text_q.device) < text_kv_offsets[-1] + ) & has_caption # [S_text] + text_res = torch.where(real_text_rows[None, :, None, None], text_res, 0) # [1,S_text,H,D] text_out = text_res.squeeze(0).flatten(-2, -1) return from_mode_splits(text_out, video_out, packed_query_states) @@ -2250,6 +2268,7 @@ def dispatch_attention_with_memory( memory_v=memory_v, text_sink_layer=memory_value.text_sink_layer if isinstance(memory_value, MultiviewARMemoryValue) else None, prompt_layout=memory_value.prompt_layout if isinstance(memory_value, MultiviewARMemoryValue) else None, + text_cache=memory_value.text_cache, ) return output, None if isinstance(memory_value, KVTrainMemoryValue): diff --git a/cosmos_framework/model/generator/mot/causal_context_parallel_test.py b/cosmos_framework/model/generator/mot/causal_context_parallel_test.py index 4681bbba9..1fdfbcbf7 100644 --- a/cosmos_framework/model/generator/mot/causal_context_parallel_test.py +++ b/cosmos_framework/model/generator/mot/causal_context_parallel_test.py @@ -7,10 +7,6 @@ import torch import torch.distributed as dist -from cosmos_framework.data.generator.sequence_packing.runtime import ( - SequencePack, - get_gen_seq, -) from cosmos_framework.model.attention import attention as full_attention from cosmos_framework.model.generator.mot.attention import SplitInfo from cosmos_framework.model.generator.mot.causal_attention import ( @@ -22,6 +18,10 @@ ) from cosmos_framework.model.generator.mot.parallelize_unified_mot import ContextParallelDispatch from cosmos_framework.model.generator.utils.kv_cache import ARMemoryValue, DualKVCache, TFNoisyMemoryValue +from cosmos_framework.data.generator.sequence_packing.runtime import ( + SequencePack, + get_gen_seq, +) from cosmos_framework.utils.generator.parallelism import ParallelDims @@ -1056,11 +1056,11 @@ def _cp_attention_ar_gen_static( cu_seqlens_kv_t = torch.tensor([0, real_total_kv_len], device=device, dtype=torch.int32) real_gen_cache_len_t = torch.tensor([real_len], device=device, dtype=torch.int32) max_seqlen_KV = s_und + gen_len + max_gen_cache_tokens - # Composite ``[und | hist | curr | pad]`` buffer; attention places ``curr`` after real history. + # Composite ``[und | curr | hist | pad]`` buffer; the attention writes ``curr`` in place. curr_room_k = gen_k_buf.new_zeros((1, gen_len, gen_k_buf.shape[2], gen_k_buf.shape[3])) curr_room_v = gen_v_buf.new_zeros((1, gen_len, gen_v_buf.shape[2], gen_v_buf.shape[3])) - kv_k_static = torch.cat([und_k_cached, gen_k_buf, curr_room_k], dim=1).contiguous() - kv_v_static = torch.cat([und_v_cached, gen_v_buf, curr_room_v], dim=1).contiguous() + kv_k_static = torch.cat([und_k_cached, curr_room_k, gen_k_buf], dim=1).contiguous() + kv_v_static = torch.cat([und_v_cached, curr_room_v, gen_v_buf], dim=1).contiguous() mv = ARMemoryValue( und_k_cached=None, @@ -1074,7 +1074,7 @@ def _cp_attention_ar_gen_static( kv_k_static=kv_k_static, kv_v_static=kv_v_static, static_curr_offset=s_und, - static_hist_offset=s_und, + static_hist_offset=s_und + gen_len, real_gen_cache_len_t=real_gen_cache_len_t, cu_seqlens_q_t=cu_seqlens_q_t, cu_seqlens_kv_t=cu_seqlens_kv_t, diff --git a/cosmos_framework/model/generator/mot/causal_cosmos3_vfm_network.py b/cosmos_framework/model/generator/mot/causal_cosmos3_vfm_network.py index 680e47078..bf3d13810 100644 --- a/cosmos_framework/model/generator/mot/causal_cosmos3_vfm_network.py +++ b/cosmos_framework/model/generator/mot/causal_cosmos3_vfm_network.py @@ -289,6 +289,8 @@ def _replace_teacher_forcing_mask( ) flex_metadata = replace(flex_metadata, teacher_forcing_replay_policy=teacher_forcing_replay_policy) setattr(attention_meta, "replay_maskless_plan", build_replay_maskless_plan(flex_metadata, sensor_items)) + if rolling_request is not None: + setattr(attention_meta, "ar_text_cache_layers", rolling_request.text_cache_layers) # The marked replay dispatch must consume this plan, never the base # bidirectional plan that was constructed before the pre-hook. attention_meta.multiview_maskless = None diff --git a/cosmos_framework/model/generator/mot/context_parallel_test.py b/cosmos_framework/model/generator/mot/context_parallel_test.py index c61a6272d..c9e4cd2dd 100644 --- a/cosmos_framework/model/generator/mot/context_parallel_test.py +++ b/cosmos_framework/model/generator/mot/context_parallel_test.py @@ -19,6 +19,7 @@ from cosmos_framework.model.generator.algorithm.loss.load_balancing import compute_load_balancing_loss from cosmos_framework.configs.base.defaults.multiview_attention import CrossViewTemporalWindowConfig from cosmos_framework.data.generator.joint_dataloader import IterativeJointDataLoader +from cosmos_framework.model.generator.mot import context_parallel_utils as cp_utils from cosmos_framework.model.generator.mot.attention import ( SplitInfo, build_packed_sequence, @@ -31,7 +32,7 @@ ) from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean from cosmos_framework.model.generator.utils.load_balancing_stats import LBLMetadata, compute_sample_lbl_stats -from cosmos_framework.model.generator.utils.memory import MemoryValue +from cosmos_framework.model.generator.utils.memory import ARTextLayerCache, MemoryValue from cosmos_framework.data.generator.sequence_packing import ( PackedSequence, build_sequence_plans_from_data_batch, @@ -72,6 +73,7 @@ def __next__(self) -> dict[str, Any]: def _make_window_trainer(cp_size: int = 2) -> tuple[ImaginaireTrainer, Any]: trainer = object.__new__(ImaginaireTrainer) + trainer.config = SimpleNamespace(trainer=SimpleNamespace(skip_model_update=False)) trainer._cp_data_window = ContextParallelDataWindow() cp_mesh = SimpleNamespace(size=lambda: cp_size, get_group=lambda: object()) parallel_dims = SimpleNamespace(cp_enabled=True, cp_mesh=cp_mesh) @@ -1562,3 +1564,76 @@ def test_context_parallel_multiview_maskless() -> None: test_context_parallel_attention_two_way() test_get_context_parallel_sharded_sequence_three_way() test_context_parallel_multiview_maskless() + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("cp_size", [4, 16]) +def test_ar_text_cache_skips_text_collectives_and_preserves_cp_kv( + monkeypatch: pytest.MonkeyPatch, + cp_size: int, +) -> None: + """Later steps exchange media only, including CP greater than KV-head count.""" + calls: list[str] = [] + monkeypatch.setattr(dist, "get_world_size", lambda group: cp_size) + + def gather(tensor: torch.Tensor, **kwargs: Any) -> torch.Tensor: # [S_local,H,D] -> [S,H_local,D] + calls.append("gather") + return tensor.repeat(cp_size, 1, 1)[:, : tensor.shape[1] // cp_size].contiguous() # [S,H_local,D] + + def scatter(tensor: torch.Tensor, **kwargs: Any) -> torch.Tensor: # [S,H_local*D] -> [S_local,H*D] + calls.append("scatter") + return tensor[: tensor.shape[0] // cp_size].repeat(1, cp_size) # [S_local,H*D] + + monkeypatch.setattr(cp_utils, "gather_seq_scatter_heads", gather) + monkeypatch.setattr(cp_utils, "gather_heads_scatter_seq", scatter) + meta: SequencePack = { + "is_sharded": True, + "_num_causal_tokens": 2 * cp_size, + "_num_full_tokens": 3 * cp_size, + "sample_offsets": torch.tensor([0, 5 * cp_size], dtype=torch.int32), # [B+1] + "max_sample_len": 5 * cp_size, + "max_causal_len": 2 * cp_size, + "max_full_len": 3 * cp_size, + "_causal_indices": torch.arange(2 * cp_size), # [S_text] + "_full_indices": torch.arange(2 * cp_size, 5 * cp_size), # [S_gen] + "_causal_seq_offsets": torch.tensor([0, 2 * cp_size], dtype=torch.int32), # [B+1] + "_full_only_seq_offsets": torch.tensor([0, 3 * cp_size], dtype=torch.int32), # [B+1] + } + text = torch.randn(2, 32, 4) # [S_text_local,H,D] + media = torch.randn(3, 32, 4) # [S_gen_local,H,D] + q = from_mode_splits(text, media, meta, is_sharded=True) + k = from_mode_splits(text[:, :8], media[:, :8], meta, is_sharded=True) # streams [S_local,H_kv,D] + v = from_mode_splits(text[:, :8] * 2, media[:, :8] * 2, meta, is_sharded=True) # [S_local,H_kv,D] + normalized = from_mode_splits(text[:, :8] * 3, media[:, :8], meta, is_sharded=True) # [S_local,H_kv,D] + + def attention( + q: SequencePack, k: SequencePack, v: SequencePack, mask: object, **kwargs: Any + ) -> tuple[SequencePack, None]: + normalized_k = kwargs["packed_key_states_normalized"] + und = get_und_seq(q) + get_und_seq(k).repeat_interleave(4 if cp_size == 4 else 2, dim=1) # [S_text,H_local,D] + gen = get_gen_seq(q) + get_gen_seq(normalized_k).repeat_interleave( + 4 if cp_size == 4 else 2, dim=1 + ) # [S_gen,H_local,D] + return from_mode_splits(und.flatten(-2), gen.flatten(-2), q), None + + memory = MemoryValue(text_cache=ARTextLayerCache()) + mesh = SimpleNamespace(get_group=lambda: None) + with torch.no_grad(): + cp_utils.context_parallel_attention( + mesh, q, k, v, object(), attention, memory_value=memory, packed_key_states_normalized=normalized + ) + assert len(calls) == 9 + calls.clear() + cached, cached_kv = cp_utils.context_parallel_attention( + mesh, q, k, v, object(), attention, memory_value=memory, packed_key_states_normalized=normalized + ) + assert calls == ["gather", "gather", "gather", "scatter"] + reference, reference_kv = cp_utils.context_parallel_attention( + mesh, q, k, v, object(), attention, memory_value=MemoryValue(), packed_key_states_normalized=normalized + ) + torch.testing.assert_close(get_gen_seq(cached), get_gen_seq(reference), rtol=0, atol=0) + torch.testing.assert_close(get_und_seq(cached), get_und_seq(reference), rtol=0, atol=0) + assert cached_kv is not None and reference_kv is not None + for a, b in zip(cached_kv, reference_kv, strict=True): + torch.testing.assert_close(a, b, rtol=0, atol=0) diff --git a/cosmos_framework/model/generator/mot/context_parallel_utils.py b/cosmos_framework/model/generator/mot/context_parallel_utils.py index a7c6194e5..0ee459300 100644 --- a/cosmos_framework/model/generator/mot/context_parallel_utils.py +++ b/cosmos_framework/model/generator/mot/context_parallel_utils.py @@ -491,6 +491,9 @@ def context_parallel_attention( cp_group = cp_mesh.get_group() cp_world_size = torch.distributed.get_world_size(cp_group) assert cp_world_size > 1, "Context parallel world size must be greater than 1" + text_cache = memory_value.text_cache if memory_value is not None else None + exchanged_text = text_cache.exchanged if text_cache is not None else None + cached_text = exchanged_text is not None # The streams only. These packs are context parallel local shards, and the tower offsets # ``get_causal_seq`` would hand out alongside them describe the unsharded stream they were cut # from. Nothing below wants them: this pass rebuilds its own ranges from the gathered lengths. @@ -539,30 +542,43 @@ def context_parallel_attention( # when doing AR-inference with a KV-cache. if kv_head_repeats > 1: - k_und_seq = _repeat_kv_heads_for_cp(k_und_seq, kv_head_repeats) # [text_shard_len,repeated_kv_heads,head_dim] + if not cached_text: + k_und_seq = _repeat_kv_heads_for_cp( + k_und_seq, kv_head_repeats + ) # [text_shard_len,repeated_kv_heads,head_dim] + v_und_seq = _repeat_kv_heads_for_cp( + v_und_seq, kv_head_repeats + ) # [text_shard_len,repeated_kv_heads,head_dim] k_gen_seq = _repeat_kv_heads_for_cp(k_gen_seq, kv_head_repeats) # [gen_shard_len,repeated_kv_heads,head_dim] - v_und_seq = _repeat_kv_heads_for_cp(v_und_seq, kv_head_repeats) # [text_shard_len,repeated_kv_heads,head_dim] v_gen_seq = _repeat_kv_heads_for_cp(v_gen_seq, kv_head_repeats) # [gen_shard_len,repeated_kv_heads,head_dim] # all2all: gather sequence, scatter heads → head-sharded - q_und_seq = gather_seq_scatter_heads( - q_und_seq, seq_dim=0, head_dim=1, cp_mesh=cp_mesh + q_und_seq = ( + exchanged_text[0] + if exchanged_text is not None + else gather_seq_scatter_heads(q_und_seq, seq_dim=0, head_dim=1, cp_mesh=cp_mesh) ) # [text_len,H_local,head_dim] q_gen_seq = gather_seq_scatter_heads( q_gen_seq, seq_dim=0, head_dim=1, cp_mesh=cp_mesh ) # [gen_len,H_local,head_dim] - k_und_seq = gather_seq_scatter_heads( - k_und_seq, seq_dim=0, head_dim=1, cp_mesh=cp_mesh + k_und_seq = ( + exchanged_text[1] + if exchanged_text is not None + else gather_seq_scatter_heads(k_und_seq, seq_dim=0, head_dim=1, cp_mesh=cp_mesh) ) # [text_len,H_kv_local,head_dim] k_gen_seq = gather_seq_scatter_heads( k_gen_seq, seq_dim=0, head_dim=1, cp_mesh=cp_mesh ) # [gen_len,H_kv_local,head_dim] - v_und_seq = gather_seq_scatter_heads( - v_und_seq, seq_dim=0, head_dim=1, cp_mesh=cp_mesh + v_und_seq = ( + exchanged_text[2] + if exchanged_text is not None + else gather_seq_scatter_heads(v_und_seq, seq_dim=0, head_dim=1, cp_mesh=cp_mesh) ) # [text_len,H_kv_local,head_dim] v_gen_seq = gather_seq_scatter_heads( v_gen_seq, seq_dim=0, head_dim=1, cp_mesh=cp_mesh ) # [gen_len,H_kv_local,head_dim] + if text_cache is not None and not cached_text: + text_cache.exchanged = (q_und_seq, k_und_seq, v_und_seq) # Build head-sharded kv_to_store when memory is active. kv_to_store: KVToStore | None = None @@ -599,11 +615,17 @@ def context_parallel_attention( packed_key_states_normalized_: SequencePack | None = None if packed_key_states_normalized is not None: k_und_normalized_seq = packed_key_states_normalized["causal_seq"] # [text_shard_len,H_kv,head_dim] - if kv_head_repeats > 1: + if kv_head_repeats > 1 and not cached_text: k_und_normalized_seq = _repeat_kv_heads_for_cp(k_und_normalized_seq, kv_head_repeats) - k_und_normalized_seq = gather_seq_scatter_heads( - k_und_normalized_seq, seq_dim=0, head_dim=1, cp_mesh=cp_mesh - ) # [text_len,H_kv_local,head_dim] + if cached_text: + assert text_cache is not None and text_cache.exchanged_normalized_k is not None + k_und_normalized_seq = text_cache.exchanged_normalized_k # [text_len,H_kv_local,head_dim] + else: + k_und_normalized_seq = gather_seq_scatter_heads( + k_und_normalized_seq, seq_dim=0, head_dim=1, cp_mesh=cp_mesh + ) # [text_len,H_kv_local,head_dim] + if text_cache is not None: + text_cache.exchanged_normalized_k = k_und_normalized_seq packed_key_states_normalized_ = from_mode_splits(k_und_normalized_seq, k_gen_seq, meta, is_sharded=False) # dispatch_attention returns (output, kv_to_store | None) @@ -630,12 +652,18 @@ def context_parallel_attention( attn_output_gen_hp = attn_output_gen_hp[:q_gen_seq_len].contiguous() # [gen_len,H_local,head_dim] # all2all: gather heads, scatter seq → seq-sharded - attn_output_und_sp = gather_heads_scatter_seq( - attn_output_und_hp, - seq_dim=0, - head_dim=1, - cp_mesh=cp_mesh, - ) # [text_shard_len,H,head_dim] + if cached_text: + assert text_cache is not None and text_cache.scattered_output is not None + attn_output_und_sp = text_cache.scattered_output # [text_shard_len,H,head_dim] + else: + attn_output_und_sp = gather_heads_scatter_seq( + attn_output_und_hp, + seq_dim=0, + head_dim=1, + cp_mesh=cp_mesh, + ) # [text_shard_len,H,head_dim] + if text_cache is not None: + text_cache.scattered_output = attn_output_und_sp attn_output_gen_sp = gather_heads_scatter_seq( attn_output_gen_hp, seq_dim=0, diff --git a/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py b/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py index 35b10eb00..6f8fe24cd 100644 --- a/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py +++ b/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py @@ -14,7 +14,6 @@ from cosmos_framework.utils import log from cosmos_framework.configs.base.defaults.joint_attention import packing_layout from cosmos_framework.configs.base.defaults.multiview_attention import ( - CaptionAccess, CrossViewTemporalWindowConfig, MultiviewAttentionConfig, ResolvedBackend, @@ -46,7 +45,6 @@ ) from cosmos_framework.model.generator.mot.modeling_utils import TimestepEmbedder, has_noisy_tokens from cosmos_framework.model.generator.mot.multiview_attention import ( - is_single_camera_pack, reject_mixed_caption_layouts, reject_samples_reading_no_caption, resolve_multiview_backend, @@ -178,6 +176,8 @@ def __init__( # window, the folds read all of it through ``maskless_unavailable_reason``, and a copy of # each on this config could disagree with the other. self.multiview_attention_config = multiview_attention_config or MultiviewAttentionConfig() + # Retain the exported-config flag for compatibility and configuration validation. + # Runtime attention reads the explicit item layout recorded by the packer. self.multiview_action_conditioning = multiview_action_conditioning # Exported checkpoint configs can arrive as plain DictConfig objects, bypassing # the attrs converter. Bring a config saved before the current fields or window keys @@ -190,12 +190,17 @@ def __init__( getattr(attention_config, "geometry_position_encoding", "baseline") != "baseline" or getattr(attention_config, "maskless_cross_view_band_radius", 0) != 0 or getattr(attention_config, "maskless_cross_view_include_frame_zero", False) + or getattr(attention_config, "maskless_cross_view_attends_action", False) + or getattr(attention_config, "maskless_cross_view_action_attends_action", False) ): - # Action conditioning builds its own maskless plan, which carries no camera geometry - # and no cross-view band, so these settings would otherwise be silently ignored. + # Separate-camera action rows participate in cross-view attention. Geometry and bands + # need layout-specific support before these settings can be used with those rows. + # Those rows are sensor items, which already query and key the cross-view pass, so + # the camera-pose action flags have nothing to add there. raise ValueError( - "Multiview action conditioning supports neither RigRoPE/PRoPE geometry nor cross-view bands; " - "keep geometry_position_encoding='baseline' and the default maskless_cross_view_band_* settings." + "Multiview action conditioning supports neither RigRoPE/PRoPE geometry, cross-view bands nor " + "the maskless_cross_view_*action* flags; keep geometry_position_encoding='baseline' and the " + "default maskless_cross_view_* settings." ) self.temporal_compression_factor_vision = temporal_compression_factor_vision self.natten_parameter_list = natten_parameter_list @@ -343,7 +348,10 @@ def __init__(self, language_model, config: Cosmos3VFMNetworkConfig): f"control_attends_sensor={multiview.mask.control_attends_sensor}, " f"cross_view_temporal_window.bounds={multiview.mask.cross_view_temporal_window.bounds}, " "cross_view_temporal_window.includes_first_frame=" - f"{multiview.mask.cross_view_temporal_window.includes_first_frame}. It builds " + f"{multiview.mask.cross_view_temporal_window.includes_first_frame}, " + f"maskless_cross_view_attends_action={multiview.maskless_cross_view_attends_action}, " + "maskless_cross_view_action_attends_action=" + f"{multiview.maskless_cross_view_action_attends_action}. It builds " "no mask, so it imposes no alignment on the GEN stream: its partitions cover " "whatever padding the pack has." ) @@ -1476,153 +1484,43 @@ def _prepare_multiview_attention( if self.multiview_backend is None: return - mask_config = self.config.multiview_attention_config.mask - same_view_window = mask_config.same_view_temporal_window - has_same_view_window = any( - bounds is not None - for bounds in ( - same_view_window.sensor_to_sensor_bounds, - same_view_window.sensor_to_control_bounds, - same_view_window.control_to_control_bounds, - same_view_window.control_to_sensor_bounds, - ) - ) - action_conditioning = getattr(self.config, "multiview_action_conditioning", False) - counts = packed_seq.num_vision_items_per_sample - if action_conditioning and counts is None and not has_same_view_window: - # Preserve dense attention for the unchanged single-view streams - # mixed into an opt-in run; their existing action block is unchanged. - # Configured same-view windows instead need the ordinary geometry below. - return - if action_conditioning and counts is not None: - if ( - self.multiview_backend != "maskless" - or packed_seq.num_action_tokens_per_supertoken != 0 - or packed_seq.vision is None - or packed_seq.action is None - or packed_seq.sound is not None - or packed_seq.lidar is not None - or packed_seq.radar is not None - ): - raise ValueError("Multiview action conditioning requires RGB/action blocks and maskless attention.") - config = self.config.multiview_attention_config - full_seq, _ = get_full_only_seq(input_pack) - shapes: list[tuple[int, int, int]] = [] - rates: list[float] = [] - view_axis: list[int] = [] - frame_offsets: list[int] = [] - items_per_sample: list[int] = [] - factor = self.config.temporal_compression_factor_vision - if len(packed_seq.action.token_shapes) != len(counts): - raise ValueError("Multiview packs require one camera-major action payload per sample.") - vision_cursor = 0 - for sample, num_views in enumerate(counts): - view_shapes = packed_seq.vision.token_shapes[vision_cursor : vision_cursor + num_views] - view_rates = packed_seq.vision.seconds_per_frame[vision_cursor : vision_cursor + num_views] - latent_t = view_shapes[0][0] - if any(shape[0] != latent_t for shape in view_shapes) or packed_seq.action.token_shapes[sample] != ( - num_views * (latent_t - 1) * factor, - ): - raise ValueError("Multiview action blocks must retain every frame-rate action for every view.") - shapes.extend((t, h, w) for t, h, w in view_shapes) - rates.extend(view_rates) - view_axis.extend(range(num_views)) - frame_offsets.extend([0] * num_views) - # The action block stays separate. Describe its full frame-rate - # rows in latent-sized groups, sharing each RGB view's attention - # identity. Match pack_supertokens_temporal_causal's [null, V0], - # [actions[0:factor], V1], ... alignment, but omit its null prefix. - # The first real action group therefore starts at latent frame 1. - if latent_t > 1: - shapes.extend([(latent_t - 1, 1, factor)] * num_views) - rates.extend(view_rates) - view_axis.extend(range(num_views)) - frame_offsets.extend([1] * num_views) - items_per_sample.append(num_views * (2 if latent_t > 1 else 1)) - vision_cursor += num_views - caption_access: list[CaptionAccess] = ["all_captions"] * len(shapes) - # All RGB items are synchronized targets, not Auto's prefix controls/ - # last target. Conditioned actions join the same view/time groups; - # their clean inputs do not make them separate control-image streams. - attention_meta.multiview_maskless = build_multiview_maskless_plan( - [1] * len(shapes), - shapes, - device=full_seq.device, - seconds_per_frame=rates, - frame_offsets=frame_offsets, - items_per_sample=items_per_sample, - is_control=[False] * len(shapes), - view_axis=view_axis, - caption_access=caption_access, - same_view_temporal_window=config.mask.same_view_temporal_window, - padded_gen_tokens=full_seq.shape[0], - attention_scope=config.mask.attention_scope, - cross_view_temporal_window=config.mask.cross_view_temporal_window, - ) - return - # No pathway check here: ``multiview_backend`` is non-None exactly when # ``joint_attn_implementation == "multiview"``, so the two cannot disagree. # natten_metadata_list is always None here (only the three-way packer builds it). if self.natten_parameter_list: - raise ValueError("Multiview FlexAttention and NATTEN cannot be enabled together.") + raise ValueError("Multiview attention and NATTEN cannot be enabled together.") + + # Sensor-layout metadata opts a pack into multiview attention, including one-camera + # packs whose control rules or same-view temporal windows still need a plan. LiDAR + # and radar carry their own item counts even when camera-view metadata is absent. + # The dedicated robot packer records its separate-camera action layout explicitly. + # Ordinary image/video/action/sound packs leave all layout fields unset and retain + # two-way attention within each sample, independent of the multiview control rules. + if ( + packed_seq.num_views_per_vision_item is None + and packed_seq.num_lidar_items_per_sample is None + and packed_seq.num_radar_items_per_sample is None + and packed_seq.multiview_action_temporal_factor is None + ): + return + + if packed_seq.multiview_action_temporal_factor is not None and self.multiview_backend != "maskless": + raise ValueError("Separate-camera action conditioning requires maskless multiview attention.") # The multiview mask cannot describe a sound stream: its items carry a # ``(frame, view)`` grid that a sound latent has no place in, and the mask rejects - # a pack whose GEN tokens it cannot all account for. Leaving the mask unbuilt sends - # the pack to dense attention, which is what the single-view Super streams trained - # under before they were mixed with multiview ones. - # - # For a single-camera pack with no control stream that is an equivalence rather - # than an approximation: ``same_view`` holds for every pair, so the whole sensor - # square is admitted, ``is_control`` is empty so no control rule narrows anything, - # and the sample-level caption is readable from every view. Dense within-sample - # attention is the same connectivity. - # - # Both exclusions are load-bearing. A genuinely multiview pack would lose its view - # scoping silently, and a pack carrying controls would have its control edges - # widened, since the mask confines them to their own view. + # a pack whose GEN tokens it cannot all account for. Ordinary sound packs already + # returned above; a metadata-bearing pack cannot fall back without losing its + # declared view, control or temporal restrictions. if packed_seq.sound is not None: - if has_same_view_window: - raise NotImplementedError("Same-view temporal windows are not implemented for sound packs.") - if not (packed_seq.action is None and is_single_camera_pack(packed_seq)): - raise ValueError( - "Multiview FlexAttention supports vision/LiDAR/radar/action batches, not sound. " - "Only a single-camera sound pack with no action stream falls back to dense " - "attention; this one carries " - + ("an action stream." if packed_seq.action is not None else "more than one camera.") - ) - return - - # Under the maskless backend a single-camera pack is one more exact dense case. Every - # sample then owns one same-view group, and the plan slices it out of the cross-instant - # fold because its instant groups would be subsets of that group, so the folds reduce to - # attention within the sample. A control item -- an action item, or a camera item that - # conditions the next one on its stream -- joins the camera's only group, which gives its - # queries the camera rows' connectivity only while ``control_attends_sensor`` is set; off - # it, the folds withhold the target's tokens from a control query and dense would not. - # A single-camera pack has no LiDAR or radar stream, so those two are the only controls. - # Skipping the fold and its gen->und merge matters for the image and single-view Action - # streams of a mixed recipe, while true multiview packs still take the folds below. - # Same-view temporal windows also require the folds: dense attention would admit - # frames outside the configured window even when the pack has only one camera. - # Every geometry mode acts only between cameras, so a single-camera pack needs no poses. - carries_control = packed_seq.action is not None or any( - int(count) > 1 for count in packed_seq.num_vision_items_per_sample or [] - ) - if ( - self.multiview_backend == "maskless" - and is_single_camera_pack(packed_seq) - and (not carries_control or mask_config.control_attends_sensor) - and not has_same_view_window - ): - return + raise ValueError( + "Multiview attention does not support sound in packs with camera/LiDAR/radar layout metadata." + ) if packed_seq.action is not None and packed_seq.vision is None: raise ValueError("Multiview FlexAttention requires a vision generation stream for action controls.") if packed_seq.action is not None and (packed_seq.lidar is not None or packed_seq.radar is not None): raise ValueError("Multiview FlexAttention does not support joint LiDAR/radar and action-control batches.") - if packed_seq.vision is None and packed_seq.lidar is None and packed_seq.radar is None: raise ValueError("Multiview FlexAttention needs a vision or LiDAR or radar generation stream.") @@ -1677,6 +1575,10 @@ def _prepare_multiview_attention( rigrope_time_scale_s=self.config.multiview_attention_config.rigrope_time_scale_s, rigrope_head_masks=self._rigrope_head_masks(input_pack["full_only_seq"].device), rigrope_channels=self.config.multiview_attention_config.rigrope_channels, + cross_view_attends_action=self.config.multiview_attention_config.maskless_cross_view_attends_action, + cross_view_action_attends_action=( + self.config.multiview_attention_config.maskless_cross_view_action_attends_action + ), ) geometry_mode = self.config.multiview_attention_config.geometry_position_encoding if ( @@ -2075,6 +1977,8 @@ def _multiview_maskless_geometry( rigrope_time_scale_s: float = 1.0, rigrope_head_masks: tuple[torch.Tensor, torch.Tensor] | None = None, rigrope_channels: RigRopeChannels = "all", + cross_view_attends_action: bool = False, + cross_view_action_attends_action: bool = False, ) -> MultiviewMasklessPlan: """The plan :func:`~...multiview_maskless_attention.multiview_maskless_gen_attention` folds this batch by. @@ -2092,7 +1996,11 @@ def _multiview_maskless_geometry( here. ``control_attends_sensor`` is not among those: the folds express both answers, so it is passed down as a description of the attention rather than tested as a condition. - * at most one sensor item per stream per sample, beside its control item, plus at most one + * separate-camera action packs declare their layout explicitly: ragged RGB targets + and replicated full-rate actions share camera identities, with action groups starting + at latent frame one. They use the same plan builder and configuration as other packs. + Geometry and cross-view bands remain unsupported for this layout. + * otherwise, at most one sensor item per stream per sample, beside its control item, plus at most one camera-pose action control, and no sound. A camera item, a range item, a radar item, or one of each: a joint sample is served by quantising every stream's capture times onto the camera's frame grid, or comparing capture times directly when a temporal window is @@ -2163,6 +2071,12 @@ def _multiview_maskless_geometry( Own-view sensor keys are excluded from this pass. Same-view/caption attention stays intact. When includes_first_frame is enabled, also admit each sample's first frame on the other views for every query time, beside the window. + cross_view_attends_action: ``maskless_cross_view_attends_action``: the camera-pose action + control's per-view rows also key the cross-view pass, timed on the action clock from + video frame 1, where the packer lays the first action. + cross_view_action_attends_action: ``maskless_cross_view_action_attends_action``: those + rows also query the cross-view pass, keyed only by the other views' action rows, + on the same clock. Returns: The batch's plan. Asked only of a run that resolved to the folds, so there is no "did @@ -2191,130 +2105,115 @@ def _multiview_maskless_geometry( f"{len(lidar_counts)} LiDAR and {len(radar_counts)} radar item counts for " f"{num_samples} samples." + _MASKLESS_REFUSAL_TAIL ) - # The action stream's counts, resolved the way the mask resolves them: at most one control - # per sample, an item count the packed stream agrees with, and no LiDAR or radar beside it. - # In ``forward`` the mask builder has already run this on the same pack, so it is called - # again for the resolved lists rather than to re-check them, and the defaults stay in one - # place. Ahead of the per-stream counts because an action control with no camera to target - # is the more specific account of a sample that also owns no sensor item. - action_metadata = validate_multiview_pack(packed_seq) - action_counts = action_metadata.counts_per_sample - views_per_action_item = action_metadata.views_per_item - if any(count and vision_count != 1 for count, vision_count in zip(action_counts, vision_counts)): - # A transfer sample's two camera items share their view groups already, so an action - # control beside them would have two targets to condition and the plan one to give it. - raise ValueError( - f"{_MASKLESS_REFUSAL}an action control needs exactly one target vision item; " - f"vision counts are {list(vision_counts)} and action counts are {list(action_counts)}." - + _MASKLESS_REFUSAL_TAIL - ) - # At most one sensor item per stream per sample beside its control item, and at least one - # item overall. A sample owning a camera item beside a range item or a radar item is the - # joint case: the sensors run at different rates, so the plan quantises each onto the - # camera's frame grid by capture time. The *second* item on a stream is that stream's - # control item, which the folds serve; a third is an image-editing layout, which this path - # does not -- ``control_weights`` catches most of those and this catches the rest. - if any( - cam > 2 or rng > 2 or bev > 2 or cam + rng + bev < 1 - for cam, rng, bev in zip(vision_counts, lidar_counts, radar_counts) - ): - raise ValueError( - f"{_MASKLESS_REFUSAL}its per-sample item counts are vision={list(vision_counts)}, " - f"lidar={list(lidar_counts)}, radar={list(radar_counts)}. Each sample takes at most " - "one target item per stream beside that stream's control item, and at least one " - "overall; more is an image-editing layout this path does not serve." + _MASKLESS_REFUSAL_TAIL - ) + if packed_seq.multiview_action_temporal_factor is not None: + # The packer describes ragged camera targets and replicated sensor-role actions. + # Their item validation is shared with _multiview_sensor_mask_items below. + action_counts = [1] * num_samples + axes_by_offset = {view: view for view in range(max(vision_counts, default=0))} + if geometry_position_encoding != "baseline" or cross_view_band_radius != 0 or include_frame_zero: + raise ValueError("Separate-camera action conditioning does not support geometry or cross-view bands.") + if cross_view_attends_action or cross_view_action_attends_action: + raise ValueError( + "Separate-camera action rows are sensor items, which already query and key the cross-view " + "pass; the maskless_cross_view_*action* flags apply to camera-pose action controls only." + ) + else: + # The action stream's counts, resolved the way the mask resolves them: at most one control + # per sample, an item count the packed stream agrees with, and no LiDAR or radar beside it. + # In ``forward`` the mask builder has already run this on the same pack, so it is called + # again for the resolved lists rather than to re-check them, and the defaults stay in one + # place. Ahead of the per-stream counts because an action control with no camera to target + # is the more specific account of a sample that also owns no sensor item. + action_metadata = validate_multiview_pack(packed_seq) + action_counts = action_metadata.counts_per_sample + if any(count and vision_count != 1 for count, vision_count in zip(action_counts, vision_counts)): + # A transfer sample's two camera items share their view groups already, so an action + # control beside them would have two targets to condition and the plan one to give it. + raise ValueError( + f"{_MASKLESS_REFUSAL}an action control needs exactly one target vision item; " + f"vision counts are {list(vision_counts)} and action counts are {list(action_counts)}." + + _MASKLESS_REFUSAL_TAIL + ) + # At most one sensor item per stream per sample beside its control item, and at least one + # item overall. A sample owning a camera item beside a range item or a radar item is the + # joint case: the sensors run at different rates, so the plan quantises each onto the + # camera's frame grid by capture time. The *second* item on a stream is that stream's + # control item, which the folds serve; a third is an image-editing layout, which this path + # does not -- ``control_weights`` catches most of those and this catches the rest. + if any( + cam > 2 or rng > 2 or bev > 2 or cam + rng + bev < 1 + for cam, rng, bev in zip(vision_counts, lidar_counts, radar_counts) + ): + raise ValueError( + f"{_MASKLESS_REFUSAL}its per-sample item counts are vision={list(vision_counts)}, " + f"lidar={list(lidar_counts)}, radar={list(radar_counts)}. Each sample takes at most " + "one target item per stream beside that stream's control item, and at least one " + "overall; more is an image-editing layout this path does not serve." + _MASKLESS_REFUSAL_TAIL + ) - views_per_vision_item = list(packed_seq.num_views_per_vision_item or []) - if vision is not None and not views_per_vision_item: - # One view per item, as _multiview_sensor_mask_items reads the same pack: only a - # multiview dataset writes the counts, so a pack without them is single-camera by - # construction -- a single-view stream mixed into a multiview recipe. - views_per_vision_item = [1] * len(vision.token_shapes) - if vision is not None and len(views_per_vision_item) != len(vision.token_shapes): - # Camera items need the per-camera VAE metadata to say where one view's frames end. - raise ValueError( - f"{_MASKLESS_REFUSAL}it records {len(views_per_vision_item)} per-item view counts " - f"for {len(vision.token_shapes)} vision items, so there is nothing to say where one " - "view's frames end." + _MASKLESS_REFUSAL_TAIL - ) + views_per_vision_item = list(packed_seq.num_views_per_vision_item or []) + if vision is not None and not views_per_vision_item: + # One view per item, as _multiview_sensor_mask_items reads the same pack: only a + # multiview dataset writes the counts, so a pack without them is single-camera by + # construction -- a single-view stream mixed into a multiview recipe. + views_per_vision_item = [1] * len(vision.token_shapes) + if vision is not None and len(views_per_vision_item) != len(vision.token_shapes): + # Camera items need the per-camera VAE metadata to say where one view's frames end. + raise ValueError( + f"{_MASKLESS_REFUSAL}it records {len(views_per_vision_item)} per-item view counts " + f"for {len(vision.token_shapes)} vision items, so there is nothing to say where one " + "view's frames end." + _MASKLESS_REFUSAL_TAIL + ) + # Retain the established camera/LiDAR/radar axes even when a stream is absent. + # The shared mask items use global view offsets instead of sensor-axis IDs. + axes_by_offset = {0: 0} + lidar_offset = max(views_per_vision_item, default=0) + if lidar is not None: + axes_by_offset[lidar_offset] = 1 + if radar is not None: + axes_by_offset[lidar_offset + (lidar is not None)] = 2 + + # Both layouts now describe the same item contract, in packed token order. + # Shared view offsets join controls/actions to their camera and keep LiDAR/radar + # distinct. Map offsets onto the folds' opaque view axes while retaining the + # existing camera/LiDAR/radar axis IDs for the camera-major layout. + items = [item for sample_items in sensor_mask_items for item in sample_items] + num_views = [item.num_views for item in items] + token_shapes = [_mask_item_token_shape(item.token_shape) for item in items] + rates = [item.seconds_per_frame for item in items] + items_per_sample = [len(sample_items) for sample_items in sensor_mask_items] + is_control = [item.is_control for item in items] + view_axis = [axes_by_offset[item.view_offset] for item in items] + frame_offsets = [item.frame_offset for item in items] + cross_view_key_only = [False] * len(items) + cross_view_control_queries = [False] * len(items) + if cross_view_attends_action or cross_view_action_attends_action: + # A camera-pose action control is last in its sample, as the packer lays it. + item_start = action_cursor = 0 + for sample_items, action_count in zip(sensor_mask_items, action_counts): + for index in range(item_start + len(sample_items) - action_count, item_start + len(sample_items)): + assert packed_seq.action is not None + if not ( + packed_seq.action.seconds_per_frame_known + and packed_seq.action.seconds_per_frame_known[action_cursor] + ): + raise ValueError( + f"{_MASKLESS_REFUSAL}the maskless_cross_view_*action* flags time action tokens by " + f"capture time, but action item {action_cursor} carries no action fps." + _MASKLESS_REFUSAL_TAIL + ) + cross_view_key_only[index] = cross_view_attends_action + cross_view_control_queries[index] = cross_view_action_attends_action + # Row r of a camera-pose action is the step into video frame r + 1: the multiview + # dataset packs it with action_start_frame_offset=1. The offset only times its + # cross-view tokens, so it stays 0 otherwise, which same-view windows require. + frame_offsets[index] = 1 + action_cursor += 1 + item_start += len(sample_items) - num_views: list[int] = [] - token_shapes: list[tuple[int, int, int]] = [] - rates: list[float] = [] - items_per_sample: list[int] = [] - is_control: list[bool] = [] - view_axis: list[int] = [] - vision_cursor = lidar_cursor = radar_cursor = action_cursor = 0 - for vision_count, lidar_count, radar_count, action_count in zip( - vision_counts, lidar_counts, radar_counts, action_counts - ): - items_per_sample.append(vision_count + lidar_count + radar_count + action_count) - # The camera item first, which is the order the packer lays a sample down and the order - # the plan anchors on: a joint sample quantises capture time onto its *first* item's - # frame grid, so anchoring on the camera keeps its tokens on the frame indices a - # camera-only sample would give them. - # Within a stream every item but the last conditions the one after it, which is the - # convention _multiview_sensor_mask_items marks is_control by and the packer forces - # fully clean. Camera items share one view axis and range items another, so a camera's - # view 0 and a sweep are never the same view -- the plan's counterpart to the view - # offset the mask gives a range clip. - for index in range(vision_count): - assert vision is not None - num_views.append(views_per_vision_item[vision_cursor]) - token_shapes.append(tuple(vision.token_shapes[vision_cursor])) # type: ignore[arg-type] - rates.append(float(vision.seconds_per_frame[vision_cursor])) - is_control.append(index < vision_count - 1) - view_axis.append(0) - vision_cursor += 1 - for index in range(lidar_count): - assert lidar is not None - # A sweep fuses the whole rig rather than covering one of its cameras, so a range - # item is one "view" over its own grid -- the same count _multiview_sensor_mask_items - # gives it. The folds do not otherwise care which sensor produced the tokens. - num_views.append(1) - token_shapes.append(tuple(lidar.token_shapes[lidar_cursor])) # type: ignore[arg-type] - rates.append(float(lidar.seconds_per_frame[lidar_cursor])) - is_control.append(index < lidar_count - 1) - view_axis.append(1) - lidar_cursor += 1 - for index in range(radar_count): - assert radar is not None - # A BEV clip covers the scene around the rig rather than one of its cameras, so it - # is one "view", as a sweep is -- but on a third axis: a radar frame and a sweep - # are no more the same view than either is a camera's, and the two sensors cycle - # at 20 Hz and 10 Hz, so they do not even share an instant. - num_views.append(1) - token_shapes.append(tuple(radar.token_shapes[radar_cursor])) # type: ignore[arg-type] - rates.append(float(radar.seconds_per_frame[radar_cursor])) - is_control.append(index < radar_count - 1) - view_axis.append(2) - radar_cursor += 1 - for _index in range(action_count): - assert action is not None - # A trajectory per camera, laid view-outer like the camera item it conditions, so it - # sits on the cameras' axis and each view's run joins that view's group. It is last - # in its sample, as the packer lays it, so it never becomes the anchor. - num_views.append(views_per_action_item[action_cursor]) - token_shapes.append(_mask_item_token_shape(tuple(action.token_shapes[action_cursor]))) # type: ignore[arg-type] - rates.append(float(action.seconds_per_frame[action_cursor]) if action.seconds_per_frame else 1.0) - is_control.append(True) - view_axis.append(0) - action_cursor += 1 - - # The divisibility of each latent_t by its view count is checked by SensorMaskItem, which the - # caller built for these same items before reaching here, and again by the plan builder. - # Captions as (view_id, num_tokens) per sample, the two parallel lists the packer records. - # None keeps the gen->und pass on its per-sample form, which is what a single caption wants. - # Flattened in the same order this walked the items above -- per sample, its vision items - # then its LiDAR ones then its radar ones -- which is the order the packer lays a sample down - # and the order _multiview_sensor_mask_items builds in. The count check holds the two together. - caption_accesses = [item.caption_access for sample_items in sensor_mask_items for item in sample_items] - if len(caption_accesses) != len(num_views): - raise ValueError( - f"{_MASKLESS_REFUSAL}the mask describes {len(caption_accesses)} items for this batch and " - f"the folds derive {len(num_views)}; the two read the same pack and must agree." + _MASKLESS_REFUSAL_TAIL - ) + # Divisibility and condition-mask lengths were checked when the shared items + # were constructed. Their caption access is also identical for both backends. + # None keeps gen->und on its per-sample form for a single shared caption. + caption_accesses = [item.caption_access for item in items] captions = ( [ list(zip(view_ids, lens)) @@ -2328,10 +2227,13 @@ def _multiview_maskless_geometry( token_shapes, device=device, seconds_per_frame=rates, + frame_offsets=frame_offsets, items_per_sample=items_per_sample, is_control=is_control, control_attends_sensor=control_attends_sensor, view_axis=view_axis, + cross_view_key_only=cross_view_key_only, + cross_view_control_queries=cross_view_control_queries, captions=captions, attention_scope=attention_scope, cross_view_temporal_window=cross_view_temporal_window, @@ -2595,6 +2497,85 @@ def _multiview_caption_mask_items(packed_seq: PackedSequence) -> list[list[Capti ] +def _multiview_action_conditioning_items(packed_seq: PackedSequence) -> list[list[SensorMaskItem]]: + """Describe separate-camera RGB and frame-rate actions using the shared item contract.""" + vision, action = packed_seq.vision, packed_seq.action + counts = packed_seq.num_vision_items_per_sample + factor = packed_seq.multiview_action_temporal_factor + if ( + vision is None + or action is None + or counts is None + or factor is None + or factor < 1 + or packed_seq.num_action_tokens_per_supertoken != 0 + or packed_seq.sound is not None + or packed_seq.lidar is not None + or packed_seq.radar is not None + ): + raise ValueError("Separate-camera action conditioning requires RGB and full frame-rate action blocks.") + if ( + len(counts) != len(packed_seq.sample_lens) + or any(count < 1 for count in counts) + or sum(counts) != len(vision.token_shapes) + or len(action.token_shapes) != len(counts) + ): + raise ValueError("Multiview packs require separate RGB items and one camera-major action payload per sample.") + if packed_seq.num_views_per_vision_item not in (None, [1] * sum(counts)): + raise ValueError("Separate-camera RGB items must each contain exactly one view.") + + sensor_items: list[list[SensorMaskItem]] = [] + vision_cursor = 0 + for sample, num_views in enumerate(counts): + view_shapes = vision.token_shapes[vision_cursor : vision_cursor + num_views] + latent_t = view_shapes[0][0] + if any(shape[0] != latent_t for shape in view_shapes) or action.token_shapes[sample] != ( + num_views * (latent_t - 1) * factor, + ): + raise ValueError("Multiview action blocks must retain every frame-rate action for every view.") + sample_items: list[SensorMaskItem] = [] + for view, shape in enumerate(view_shapes): + index = vision_cursor + view + sample_items.append( + SensorMaskItem( + token_shape=_mask_item_token_shape(shape), + condition_mask=vision.condition_mask[index], + num_views=1, + view_offset=view, + is_control=False, + seconds_per_frame=vision.seconds_per_frame[index], + caption_access="all_captions", + ) + ) + # The action block stays separate. Describe its full frame-rate rows in + # latent-sized groups, sharing each RGB view's attention identity. Match + # pack_supertokens_temporal_causal's [null, V0], [actions[0:factor], V1], ... + # alignment, but omit its null prefix: the first real group starts at frame 1. + if latent_t > 1: + action_conditions = ( + action.condition_mask[sample].reshape(num_views, latent_t - 1, factor).all(-1) + ) # [V,T-1] + for view in range(num_views): + sample_items.append( + SensorMaskItem( + token_shape=(latent_t - 1, 1, factor), + condition_mask=action_conditions[view], + num_views=1, + view_offset=view, + is_control=False, + seconds_per_frame=vision.seconds_per_frame[vision_cursor + view], + caption_access="all_captions", + frame_offset=1, + ) + ) + # All RGB items are synchronized targets, not prefix controls/last target. + # Conditioned actions join those same view/time groups; clean inputs do not + # make them separate control-image streams. Every item reads the shared caption. + sensor_items.append(sample_items) + vision_cursor += num_views + return sensor_items + + def _multiview_sensor_mask_items( packed_seq: PackedSequence, *, @@ -2603,11 +2584,11 @@ def _multiview_sensor_mask_items( ) -> list[list[SensorMaskItem]]: """Describe each sample to the multiview mask in packed generation order. - A pack that records no per-camera view counts is read as one camera per item. Only a - multiview dataset writes those counts, via ``enable_per_camera_vae_encoding``, so their - absence says the pack is single-view rather than that a recipe forgot something -- which - is what lets a recipe mix single-view Super streams into a multiview run without - declaring it anywhere. + Separate-camera action packs explicitly record their layout and are normalized by + ``_multiview_action_conditioning_items``. For the usual camera-major layout, a pack + that records no per-camera view counts is read as one camera per item. Multiview + datasets write those counts via ``enable_per_camera_vae_encoding``; ordinary Super + streams leave them unset. The packer lays a sample out as vision items, LiDAR items, radar items, then action items, so walking the streams sample by sample reproduces the packed order the mask @@ -2659,6 +2640,9 @@ def _multiview_sensor_mask_items( run at different rates -- what ``cross_view_temporal_window.bounds`` needs to compare the two streams by actual capture time rather than by frame index. """ + if packed_seq.multiview_action_temporal_factor is not None: + return _multiview_action_conditioning_items(packed_seq) + num_samples = len(packed_seq.sample_lens) vision = packed_seq.vision lidar = packed_seq.lidar diff --git a/cosmos_framework/model/generator/mot/flex_attention.py b/cosmos_framework/model/generator/mot/flex_attention.py index 350ee5ac7..789645730 100644 --- a/cosmos_framework/model/generator/mot/flex_attention.py +++ b/cosmos_framework/model/generator/mot/flex_attention.py @@ -546,7 +546,7 @@ class SensorMaskItem: views, or whose condition mask does not cover that axis, is malformed on its own terms, and catching it here names the one item rather than an index into a flattened list. - Every field is required. None of them has a default, for the reason :class:`FlexMetadata` + The clip's fields are required, for the reason :class:`FlexMetadata` states of its own: each one changes what the mask admits, and a silently narrower or wider mask is the one failure this type cannot surface. ``num_views`` and ``view_offset`` place the item on the view grid every rule compares in, ``is_control`` moves it between two sets @@ -584,6 +584,8 @@ class SensorMaskItem: against the other sensors and its control stream alone. This is also what says which views a per-view caption layout has to cover: the cameras' and no others, see :func:`_camera_views_of_sample`. + frame_offset: first frame on the item's clock. Clips start at zero by default; + separate-camera action groups start at one, after the initial observation. """ token_shape: tuple[int, ...] @@ -593,8 +595,11 @@ class SensorMaskItem: is_control: bool seconds_per_frame: float caption_access: CaptionAccess + frame_offset: int = 0 def __post_init__(self) -> None: + if self.frame_offset < 0: + raise ValueError(f"SensorMaskItem.frame_offset must be nonnegative, got {self.frame_offset}.") if self.num_views < 1 or self.latent_t % self.num_views != 0: raise ValueError( f"SensorMaskItem has latent_t={self.latent_t}, which is not divisible by num_views={self.num_views}." @@ -961,6 +966,7 @@ def build_multiview_flex_metadata( # Frame index cycles within each view; view index is constant across a view's # whole frame run. Both expand to one entry per token: [item_tokens]. item_frame_ids = torch.arange(frames_per_view, device=device).repeat(num_views) # [num_views*frames_per_view] + item_frame_ids = item_frame_ids + item.frame_offset # [num_views*frames_per_view] frame_ids.append(item_frame_ids.repeat_interleave(spatial_tokens)) # [item_tokens] view_ids.append( (item.view_offset + torch.arange(num_views, device=device)).repeat_interleave( diff --git a/cosmos_framework/model/generator/mot/flex_attention_test.py b/cosmos_framework/model/generator/mot/flex_attention_test.py index 71114e272..a408916ec 100644 --- a/cosmos_framework/model/generator/mot/flex_attention_test.py +++ b/cosmos_framework/model/generator/mot/flex_attention_test.py @@ -21,6 +21,7 @@ CAPTION_SCOPE_NONE, CAPTION_SCOPE_SAME_VIEW, AttentionScope, + BackendPreference, CrossViewTemporalWindowConfig, MultiviewAttentionConfig, MultiviewAttentionMaskConfig, @@ -2972,21 +2973,112 @@ def test_exact_windowed_cross_view_pass_over_an_action_pack_matches_the_flex_mas assert not cross[action_rows].any() and not cross[:, action_rows].any() +def _cross_view_counts(plan) -> torch.Tensor: # [N_gen,N_gen] + cross = torch.zeros((plan.num_gen_tokens, plan.num_gen_tokens), dtype=torch.int64) + for partition in plan.cross_view_partitions: + for qs, qe, ks, ke in zip( + partition.query_offsets[:-1], + partition.query_offsets[1:], + partition.key_offsets[:-1], + partition.key_offsets[1:], + strict=True, + ): + queries = partition.query_indices[int(qs) : int(qe)] + keys = partition.key_indices[int(ks) : int(ke)] + cross[queries[:, None], keys] += 1 + return cross + + +@pytest.mark.L0 +@torch.no_grad() +@pytest.mark.parametrize("include_first_frame", [False, True]) +@pytest.mark.parametrize(("attends_action", "action_attends_action"), [(True, False), (False, True), (True, True)]) +def test_cross_view_action_tokens_join_other_views_by_capture_time( + include_first_frame: bool, attends_action: bool, action_attends_action: bool +) -> None: + """Action rows join the cross-view pass by capture time, from video frame 1, under either flag. + + ``attends_action`` makes them keys of the other views' camera queries; ``action_attends_action`` + makes them queries keyed by the other views' action rows alone. Neither changes what a camera + query reads of vision, and an action query never reads vision. No action row sits at capture + time 0, so the first-frame rule adds the other views' vision frame 0 alone and no action row. + """ + pack = _mask_items_partly_posed_action_pack(num_views=2) + # Latent frames 4 video frames apart at 30 fps; one action row per video frame. + assert pack.vision is not None and pack.action is not None + pack.vision.seconds_per_frame = [4 / 30, 4 / 30] + pack.vision.seconds_per_frame_known = [True, True] + pack.action.seconds_per_frame = [1 / 30] + pack.action.seconds_per_frame_known = [True] + # A window just over one video frame: the latent frame at 0 s reaches action row 0 (video + # frame 1), the one at 4/30 s reaches row 2 (video frame 3), and an action row reaches its + # neighbouring rows. + options = dict( + cross_view_temporal_window=CrossViewTemporalWindowConfig( + bounds=(-0.04, 0.04), includes_first_frame=include_first_frame + ), + ) + baseline = _cross_view_counts(_multiview_maskless_geometry_for_test(pack, **options)) + plan = _multiview_maskless_geometry_for_test( + pack, + cross_view_attends_action=attends_action, + cross_view_action_attends_action=action_attends_action, + **options, + ) + cross = _cross_view_counts(plan) + + assert plan.is_control == (False, False, True) + assert int(cross.max()) <= 1, "A query keys each token at most once per pass." + # Sample 1 packs vision view 0 at 4:6, view 1 at 6:8, then action rows view 0 at 8:11 and + # view 1 at 11:14. + assert torch.equal(cross[:, :8], baseline[:, :8]), "Vision keys are unchanged." + vision_to_action = torch.zeros((8, 6), dtype=torch.int64) + if attends_action: + vision_to_action[4, 3] = vision_to_action[5, 5] = vision_to_action[6, 0] = vision_to_action[7, 2] = 1 + assert torch.equal(cross[:8, 8:], vision_to_action) + rows = torch.arange(3) + near = ((rows[:, None] - rows[None, :]).abs() <= 1).long() # [3,3], rows within one video frame + action_to_action = torch.zeros((6, 6), dtype=torch.int64) + if action_attends_action: + action_to_action[:3, 3:] = action_to_action[3:, :3] = near + assert torch.equal(cross[8:, 8:], action_to_action) + assert not cross[8:, :8].any(), "Action queries never key vision in the cross-view pass." + + +@pytest.mark.L0 +@torch.no_grad() +@pytest.mark.parametrize("flag", ["cross_view_attends_action", "cross_view_action_attends_action"]) +def test_cross_view_action_flags_need_a_known_action_clock(flag: str) -> None: + pack = _mask_items_partly_posed_action_pack(num_views=2) + with pytest.raises(ValueError, match="carries no action fps"): + _multiview_maskless_geometry_for_test( + pack, + cross_view_temporal_window=CrossViewTemporalWindowConfig(bounds=(-0.67, 0.67)), + **{flag: True}, + ) + + @pytest.mark.L0 @torch.no_grad() @pytest.mark.parametrize( - "num_views,with_action,control_attends_sensor,takes_folds", + "num_views,with_action,control_attends_sensor,with_view_metadata,sensor_window", [ - # One view per item: the folds reduce to attention within the sample, so dense serves it. - pytest.param(1, False, False, False, id="single_camera"), - pytest.param(1, True, True, False, id="single_camera_action_reads_vision"), + # Metadata marks the multiview pathway even when every item has just one view. + pytest.param(1, False, False, True, None, id="single_camera"), + pytest.param(1, True, True, True, None, id="single_camera_action_reads_vision"), # Off the flag an action query may not read its camera, which dense would let it. - pytest.param(1, True, False, True, id="single_camera_action_withheld_from_vision"), - pytest.param(2, True, True, True, id="multiview"), + pytest.param(1, True, False, True, None, id="single_camera_action_withheld_from_vision"), + pytest.param(2, True, True, True, None, id="multiview"), + pytest.param(1, True, False, False, None, id="untagged_action_keeps_two_way"), + pytest.param(1, False, False, True, (-1, 0), id="single_camera_same_view_window"), ], ) -def test_maskless_runs_send_only_exactly_dense_packs_to_dense_attention( - num_views: int, with_action: bool, control_attends_sensor: bool, takes_folds: bool +def test_multiview_routing_uses_camera_metadata( + num_views: int, + with_action: bool, + control_attends_sensor: bool, + with_view_metadata: bool, + sensor_window: tuple[int, int] | None, ) -> None: pytest.importorskip("transformers", reason="cosmos3_vfm_network requires the Cosmos3 network dependencies.") from cosmos_framework.model.generator.mot.cosmos3_vfm_network import Cosmos3VFMNetwork @@ -2997,6 +3089,9 @@ def test_maskless_runs_send_only_exactly_dense_packs_to_dense_attention( pack.num_views_per_action_item = None pack.num_action_items_per_sample = None pack.sample_lens = [2 * num_views, 2 * num_views] + if not with_view_metadata: + pack.num_views_per_vision_item = None + pack.num_views_per_action_item = None network = SimpleNamespace( multiview_backend="maskless", natten_parameter_list=None, @@ -3008,6 +3103,7 @@ def test_maskless_runs_send_only_exactly_dense_packs_to_dense_attention( mask=MultiviewAttentionMaskConfig( attention_scope="decomposed", control_attends_sensor=control_attends_sensor, + same_view_temporal_window=SameViewTemporalWindowConfig(sensor_to_sensor_bounds=sensor_window), ), ) ), @@ -3018,31 +3114,120 @@ def test_maskless_runs_send_only_exactly_dense_packs_to_dense_attention( if modality is not None for token_shape in modality.token_shapes ) - attention_meta = SimpleNamespace(multiview_maskless=None) + attention_meta = SimpleNamespace(multiview_maskless=None, flex_block_mask=None) + full_only_seq = torch.zeros(num_gen_tokens, 1) # [N_gen,1] Cosmos3VFMNetwork._prepare_multiview_attention( cast(Cosmos3VFMNetwork, network), pack, - cast(SequencePack, {"full_only_seq": torch.zeros(num_gen_tokens, 1)}), + cast(SequencePack, {"full_only_seq": full_only_seq}), cast(SplitInfo, attention_meta), ) - assert (attention_meta.multiview_maskless is not None) is takes_folds + plan = attention_meta.multiview_maskless + assert (plan is not None) is with_view_metadata + if plan is not None: + assert plan.control_attends_sensor is control_attends_sensor + assert plan.sensor_to_sensor_window == sensor_window + if with_action: + assert plan.is_control[-1] + if sensor_window is not None: + assert plan.neighborhood_layout == (2, 1, 1) + assert plan.sensor_gather.tolist() == [0, 1, 2, 3] + else: + assert attention_meta.flex_block_mask is None @pytest.mark.L0 @torch.no_grad() @pytest.mark.parametrize( - "control_attends_sensor,takes_folds", - [pytest.param(True, False, id="dense"), pytest.param(False, True, id="folds")], + "backend,num_views,control_stream,with_view_metadata,control_attends_sensor", + [ + pytest.param("maskless", 1, "action", False, False, id="untagged_sound_action_maskless"), + pytest.param("flex_triton", 1, "action", False, False, id="untagged_sound_action_flex"), + pytest.param("maskless", 1, "vision", False, False, id="untagged_sound_vision_control"), + pytest.param("maskless", 1, "none", False, False, id="untagged_plain_sound_maskless"), + pytest.param("flex_triton", 1, "none", False, False, id="untagged_plain_sound_flex"), + pytest.param("maskless", 1, "none", True, True, id="tagged_single_camera_sound_maskless"), + pytest.param("flex_triton", 1, "none", True, False, id="tagged_single_camera_sound_flex"), + pytest.param("maskless", 2, "none", True, True, id="tagged_multiview_sound"), + ], ) -def test_maskless_per_camera_prope_single_camera_packs_need_no_poses( - control_attends_sensor: bool, takes_folds: bool +def test_multiview_sound_routing_uses_camera_metadata( + backend: BackendPreference, + num_views: int, + control_stream: str, + with_view_metadata: bool, + control_attends_sensor: bool, ) -> None: + """A standard batch can pack sound and actions from different single-camera samples. + + The sound stream belongs to the first sample and any action stream to the second. + Untagged samples retain their original two-way attention regardless of multiview + control settings. Any pack tagged with sensor metadata must reject sound. + """ + pytest.importorskip("transformers", reason="cosmos3_vfm_network requires the Cosmos3 network dependencies.") + from cosmos_framework.model.generator.mot.cosmos3_vfm_network import Cosmos3VFMNetwork + + if control_stream == "vision": + pack = _mask_items_pack(num_views=num_views, with_lidar=False) + pack.sample_lens = [4 * num_views] + else: + pack = _mask_items_partly_posed_action_pack(num_views=num_views) + if control_stream == "none": + pack.action = None + pack.num_action_items_per_sample = None + pack.num_views_per_action_item = None + pack.sample_lens = [2 * num_views, 2 * num_views] + if not with_view_metadata: + # Ordinary single-view streams do not carry per-camera view metadata. + pack.num_views_per_vision_item = None + pack.num_views_per_action_item = None + pack.sound = ModalityData( + tokens=[torch.zeros(1, 3)], # list[[C,T_sound]] + token_shapes=[(3, 1, 1)], + condition_mask=[torch.zeros(3)], # list[[T_sound]] + seconds_per_frame=[1.0], + ) + pack.sample_lens[0] += 3 + pack.modality_sample_ids["sound"] = [0] + network = SimpleNamespace( + multiview_backend=backend, + natten_parameter_list=None, + config=SimpleNamespace( + multiview_attention_config=MultiviewAttentionConfig( + backend=backend, + mask=MultiviewAttentionMaskConfig( + attention_scope="decomposed", + control_attends_sensor=control_attends_sensor, + ), + ) + ), + ) + attention_meta = SimpleNamespace(multiview_maskless=None, flex_block_mask=None) + full_only_seq = torch.zeros(sum(pack.sample_lens), 1) # [N_gen,1] + expected = pytest.raises(ValueError, match="sound") if with_view_metadata else contextlib.nullcontext() + + with expected: + Cosmos3VFMNetwork._prepare_multiview_attention( + cast(Cosmos3VFMNetwork, network), + pack, + cast(SequencePack, {"full_only_seq": full_only_seq}), + cast(SplitInfo, attention_meta), + ) + + assert attention_meta.multiview_maskless is None + assert attention_meta.flex_block_mask is None + + +@pytest.mark.L0 +@torch.no_grad() +@pytest.mark.parametrize("control_attends_sensor", [True, False]) +def test_maskless_per_camera_prope_single_camera_packs_need_no_poses(control_attends_sensor: bool) -> None: """Per-camera PRoPE acts only between cameras, so a single-camera pack carries no poses. - The model gives single-camera items no poses. Such a pack takes the dense shortcut when - it may, and otherwise the folds, without failing the multiview geometry check. + The model gives single-camera items no poses. Metadata still routes such a pack through + the folds without failing the multiview geometry check. """ pytest.importorskip("transformers", reason="cosmos3_vfm_network requires the Cosmos3 network dependencies.") from cosmos_framework.model.generator.mot.cosmos3_vfm_network import Cosmos3VFMNetwork @@ -3076,33 +3261,33 @@ def test_maskless_per_camera_prope_single_camera_packs_need_no_poses( ), ) attention_meta = SimpleNamespace(multiview_maskless=None) + full_only_seq = torch.zeros(2 * 6 * 2 * 3, 1) # [N_gen,1] Cosmos3VFMNetwork._prepare_multiview_attention( cast(Cosmos3VFMNetwork, network), pack, - cast(SequencePack, {"full_only_seq": torch.zeros(2 * 6 * 2 * 3, 1)}), + cast(SequencePack, {"full_only_seq": full_only_seq}), cast(SplitInfo, attention_meta), ) - assert (attention_meta.multiview_maskless is not None) is takes_folds - if takes_folds: - assert attention_meta.multiview_maskless.camera_relative_pose is None + assert attention_meta.multiview_maskless is not None + assert attention_meta.multiview_maskless.camera_relative_pose is None @pytest.mark.L0 @torch.no_grad() @pytest.mark.parametrize( - "control_attends_sensor,sensor_window,takes_folds", + "control_attends_sensor,sensor_window", [ - pytest.param(True, None, False, id="control_reads_target"), + pytest.param(True, None, id="control_reads_target"), # Off the flag the camera control may not read its target, which dense would let it. - pytest.param(False, None, True, id="control_withheld_from_target"), + pytest.param(False, None, id="control_withheld_from_target"), # Even with controls readable, a causal sensor window excludes future target frames. - pytest.param(True, (-2, 0), True, id="causal_sensor_window"), + pytest.param(True, (-2, 0), id="causal_sensor_window"), ], ) -def test_maskless_single_camera_transfer_pack_goes_dense_only_when_its_control_reads_the_target( - control_attends_sensor: bool, sensor_window: tuple[int, int] | None, takes_folds: bool +def test_maskless_single_camera_transfer_pack_preserves_control_visibility( + control_attends_sensor: bool, sensor_window: tuple[int, int] | None ) -> None: """A camera item conditioning the next one is a control even without an action stream.""" pytest.importorskip("transformers", reason="cosmos3_vfm_network requires the Cosmos3 network dependencies.") @@ -3136,17 +3321,62 @@ def test_maskless_single_camera_transfer_pack_goes_dense_only_when_its_control_r ), ) attention_meta = SimpleNamespace(multiview_maskless=None) + full_only_seq = torch.zeros(2 * 6 * 2 * 3, 1) # [N_gen,1] + + Cosmos3VFMNetwork._prepare_multiview_attention( + cast(Cosmos3VFMNetwork, network), + pack, + cast(SequencePack, {"full_only_seq": full_only_seq}), + cast(SplitInfo, attention_meta), + ) + + assert attention_meta.multiview_maskless is not None + assert attention_meta.multiview_maskless.control_attends_sensor is control_attends_sensor + assert attention_meta.multiview_maskless.is_control == (True, False) + assert attention_meta.multiview_maskless.sensor_to_sensor_window == sensor_window + + +@pytest.mark.L0 +@torch.no_grad() +@pytest.mark.parametrize("sensor", ["lidar", "radar"]) +def test_multiview_routing_uses_range_sensor_metadata_without_camera_metadata(sensor: str) -> None: + """Either range sensor's item counts opt the pack into multiview attention on their own.""" + pytest.importorskip("transformers", reason="cosmos3_vfm_network requires the Cosmos3 network dependencies.") + from cosmos_framework.model.generator.mot.cosmos3_vfm_network import Cosmos3VFMNetwork + + stream = _lidar_stream([(3, 1, 1)]) + pack = PackedSequence( + sample_lens=[3], + lidar=stream if sensor == "lidar" else None, + radar=stream if sensor == "radar" else None, + num_lidar_items_per_sample=[1] if sensor == "lidar" else None, + num_radar_items_per_sample=[1] if sensor == "radar" else None, + ) + network = SimpleNamespace( + multiview_backend="maskless", + natten_parameter_list=None, + rigrope=None, + _rigrope_head_masks=lambda device: None, + config=SimpleNamespace( + multiview_attention_config=MultiviewAttentionConfig( + backend="maskless", + mask=MultiviewAttentionMaskConfig(attention_scope="decomposed"), + ) + ), + ) + attention_meta = SimpleNamespace(multiview_maskless=None) + full_only_seq = torch.zeros(3, 1) # [N_gen,1] Cosmos3VFMNetwork._prepare_multiview_attention( cast(Cosmos3VFMNetwork, network), pack, - cast(SequencePack, {"full_only_seq": torch.zeros(2 * 6 * 2 * 3, 1)}), + cast(SequencePack, {"full_only_seq": full_only_seq}), cast(SplitInfo, attention_meta), ) - assert (attention_meta.multiview_maskless is not None) is takes_folds - if attention_meta.multiview_maskless is not None: - assert attention_meta.multiview_maskless.sensor_to_sensor_window == sensor_window + assert attention_meta.multiview_maskless is not None + assert attention_meta.multiview_maskless.num_views == (1,) + assert attention_meta.multiview_maskless.token_shapes == ((3, 1, 1),) @pytest.mark.L0 @@ -3249,11 +3479,16 @@ def _multiview_maskless_network(*, maskless_attention: bool, multiview: bool = T @pytest.mark.L0 @pytest.mark.parametrize("window", [None, (-0.4, 0.0)]) -def test_opt_in_action_conditioning_preserves_single_view_temporal_windows(window: TemporalWindow | None) -> None: +@pytest.mark.parametrize("with_view_metadata", [False, True]) +def test_opt_in_action_conditioning_routes_same_view_windows_from_pack_metadata( + window: TemporalWindow | None, with_view_metadata: bool +) -> None: network = _multiview_maskless_network(maskless_attention=True) network.config.multiview_action_conditioning = True network.config.multiview_attention_config.mask.same_view_temporal_window.sensor_to_sensor_bounds = window - pack = _multiview_maskless_pack(num_vision_items_per_sample=None, num_views_per_vision_item=None) + pack = _multiview_maskless_pack( + num_vision_items_per_sample=None, num_views_per_vision_item=[1] if with_view_metadata else None + ) assert pack.vision is not None pack.vision.seconds_per_frame = [0.2] attention_meta = SimpleNamespace(multiview_maskless=None) @@ -3265,11 +3500,11 @@ def test_opt_in_action_conditioning_preserves_single_view_temporal_windows(windo cast(SplitInfo, attention_meta), ) - if window is None: - assert attention_meta.multiview_maskless is None - else: + if with_view_metadata: assert attention_meta.multiview_maskless is not None - assert attention_meta.multiview_maskless.sensor_to_sensor_window == (-2, 0) + assert attention_meta.multiview_maskless.sensor_to_sensor_window == ((-2, 0) if window is not None else None) + else: + assert attention_meta.multiview_maskless is None @pytest.mark.L0 @@ -3282,7 +3517,7 @@ def test_single_camera_sound_pack_rejects_same_view_temporal_window() -> None: ) full_sequence = torch.zeros(6 * 2 * 3 + 2, 1) # [N_gen,D] - with pytest.raises(NotImplementedError, match="Same-view temporal windows.*sound"): + with pytest.raises(ValueError, match="does not support sound.*layout metadata"): network._prepare_multiview_attention( pack, cast(SequencePack, {"full_only_seq": full_sequence}), diff --git a/cosmos_framework/model/generator/mot/multiview_action_conditioning_attention_test.py b/cosmos_framework/model/generator/mot/multiview_action_conditioning_attention_test.py new file mode 100644 index 000000000..789b56251 --- /dev/null +++ b/cosmos_framework/model/generator/mot/multiview_action_conditioning_attention_test.py @@ -0,0 +1,244 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Real robot packs use the shared multiview attention metadata and configuration.""" + +from types import SimpleNamespace +from typing import cast + +import pytest +import torch + +from cosmos_framework.configs.base.defaults.model_config import MultiviewActionConditioningConfig +from cosmos_framework.configs.base.defaults.multiview_attention import ( + CrossViewTemporalWindowConfig, + MultiviewAttentionConfig, + MultiviewAttentionMaskConfig, + SameViewTemporalWindowConfig, +) +from cosmos_framework.model.generator.mot.attention import SplitInfo +from cosmos_framework.model.generator.mot.cosmos3_vfm_network import ( + Cosmos3VFMNetwork, + _multiview_sensor_mask_items, +) +from cosmos_framework.model.generator.mot.multiview_maskless_attention import ( + MultiviewMasklessPlan, + cross_view_passes, +) +from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean +from cosmos_framework.data.generator.sequence_packing.packers import ( + pack_multiview_action_conditioning, + replicate_multiview_actions, +) +from cosmos_framework.data.generator.sequence_packing.sequence import PackedSequence, SequencePlan + + +def _robot_pack(*, frames: tuple[int, ...] = (3,), num_views: int = 3) -> PackedSequence: + grids = [(1, 2), (1, 1), (2, 1)][:num_views] + views = [torch.zeros(1, 2, count, h, w) for count in frames for h, w in grids] # list[[1,2,T,H,W]] + actions = [ + torch.arange((count - 1) * 16, dtype=torch.float32).reshape((count - 1) * 4, 4) for count in frames + ] # list[[4*(T-1),4]] + fps = torch.full((len(frames),), 20.0) # [B] + data = GenerationDataClean( + batch_size=len(frames), + is_image_batch=False, + x0_tokens_vision=views, + x0_tokens_action=actions, + fps_vision=fps, + num_vision_items_per_sample=[num_views] * len(frames), + ) + data = replicate_multiview_actions(data, num_views=num_views, temporal_factor=4, action_dim=4) + plans = [ + SequencePlan( + has_text=True, + has_vision=True, + has_action=True, + condition_frame_indexes_vision=[0], + condition_frame_indexes_action=list(range((count - 1) * 4)), + ) + for count in frames + ] + timesteps = torch.full((len(frames),), 0.5) # [B] + return pack_multiview_action_conditioning( + plans, + [[3, 4] for _ in frames], + data, + timesteps, + config=MultiviewActionConditioningConfig( + view_code_start=2, view_codes=[[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]][:num_views] + ), + special_tokens={"eos_token_id": 0, "start_of_generation": 1, "end_of_generation": 2}, + patch_spatial=1, + temporal_factor=4, + temporal_margin=15_000, + base_fps=24.0, + ) + + +def _no_rigrope_heads(_device: torch.device) -> None: + return None + + +def _prepare( + packed: PackedSequence, + *, + mask: MultiviewAttentionMaskConfig | None = None, +) -> MultiviewMasklessPlan: + # Only the pack identifies the robot layout. No network-level action switch or + # temporal compression factor is available for the attention path to infer it from. + network = SimpleNamespace( + multiview_backend="maskless", + natten_parameter_list=None, + rigrope=None, + _rigrope_head_masks=_no_rigrope_heads, + config=SimpleNamespace( + multiview_attention_config=MultiviewAttentionConfig( + backend="maskless", + mask=mask or MultiviewAttentionMaskConfig(attention_scope="decomposed"), + ) + ), + ) + attention_meta = SplitInfo( + packed.split_lens, packed.attn_modes, packed.sample_lens, actual_len=packed.sequence_length + ) + gen_tokens = sum(length for length, mode in zip(packed.split_lens, packed.attn_modes) if mode == "full") + full_only_seq = torch.zeros(gen_tokens + 3, 1) # [N_gen+3,1] + Cosmos3VFMNetwork._prepare_multiview_attention( + cast(Cosmos3VFMNetwork, network), packed, {"full_only_seq": full_only_seq}, attention_meta + ) + assert attention_meta.flex_block_mask is None + assert attention_meta.multiview_maskless is not None + return attention_meta.multiview_maskless + + +def _cross_edges(plan: MultiviewMasklessPlan) -> set[tuple[int, int]]: + edges: set[tuple[int, int]] = set() + for cross_pass in cross_view_passes(plan): + queries = cross_pass.gather.tolist() + keys = cross_pass.kv_gather.tolist() + for qs, qe, ks, ke in zip( + cross_pass.offsets.tolist()[:-1], + cross_pass.offsets.tolist()[1:], + cross_pass.kv_offsets.tolist()[:-1], + cross_pass.kv_offsets.tolist()[1:], + strict=True, + ): + edges.update((query, key) for query in queries[qs:qe] for key in keys[ks:ke]) + return edges + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_robot_pack_uses_shared_items_and_preserves_ragged_full_rate_actions() -> None: + packed = _robot_pack() + assert packed.multiview_action_temporal_factor == 4 + assert packed.num_views_per_vision_item is None + assert packed.action is not None + assert packed.action.token_shapes == [(24,)] + assert packed.action.tokens[0][:, 0].tolist() == list(range(0, 32, 4)) * 3 + items = _multiview_sensor_mask_items(packed)[0] + assert [item.token_shape for item in items] == [(3, 1, 2), (3, 1, 1), (3, 2, 1)] + [(2, 1, 4)] * 3 + assert [item.view_offset for item in items] == [0, 1, 2, 0, 1, 2] + assert [item.frame_offset for item in items] == [0, 0, 0, 1, 1, 1] + assert all(not item.is_control and item.caption_access == "all_captions" for item in items) + assert [item.seconds_per_frame for item in items] == pytest.approx([0.2] * 6) + + plan = _prepare(packed) + assert plan.num_gen_tokens == 39 + assert plan.padded_gen_tokens == 42 + assert plan.is_control == (False,) * 6 + assert plan.same_view_offsets is not None and plan.same_view_gather is not None + assert plan.same_view_offsets.tolist() == [0, 14, 25, 39, 42] + assert plan.same_view_gather.tolist() == [ + *range(0, 6), + *range(15, 23), + *range(6, 9), + *range(23, 31), + *range(9, 15), + *range(31, 42), + ] + # Frame zero contains RGB alone. Each later instant adds all four actions per + # camera to its matching latent frame, retaining the camera-major payload order. + frame_rows = [ + [0, 1, 6, 9, 10], + [2, 3, 7, 11, 12, *range(15, 19), *range(23, 27), *range(31, 35)], + [4, 5, 8, 13, 14, *range(19, 23), *range(27, 31), *range(35, 39)], + ] + camera_rows = [ + [*range(0, 6), *range(15, 23)], + [*range(6, 9), *range(23, 31)], + [*range(9, 15), *range(31, 39)], + ] + camera_by_row = {row: camera for camera, rows in enumerate(camera_rows) for row in rows} + assert _cross_edges(plan) == { + (query, key) + for rows in frame_rows + for query in rows + for key in rows + if camera_by_row[query] != camera_by_row[key] + } + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_robot_pack_forwards_temporal_window_first_frame_rule() -> None: + packed = _robot_pack() + plain = _prepare( + packed, + mask=MultiviewAttentionMaskConfig( + attention_scope="decomposed", cross_view_temporal_window=CrossViewTemporalWindowConfig(bounds=(0.0, 0.0)) + ), + ) + widened = _prepare( + packed, + mask=MultiviewAttentionMaskConfig( + attention_scope="decomposed", + cross_view_temporal_window=CrossViewTemporalWindowConfig(bounds=(0.0, 0.0), includes_first_frame=True), + ), + ) + plain_edges = _cross_edges(plain) + widened_edges = _cross_edges(widened) + assert plain_edges < widened_edges + assert (4, 6) not in plain_edges and (4, 6) in widened_edges + assert (15, 6) not in plain_edges and (15, 6) in widened_edges + # Action[0] belongs to frame one, so adding the first observation never adds + # that action to a later query whose zero-width time window excludes it. + assert (4, 23) not in widened_edges + assert (2, 23) in plain_edges + assert (2, 15) not in widened_edges # Own-view edges remain in the same-view pass. + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_robot_pack_rejects_unsupported_same_view_window() -> None: + with pytest.raises(ValueError, match="frame offsets|neighborhood"): + _prepare( + _robot_pack(), + mask=MultiviewAttentionMaskConfig( + attention_scope="decomposed", + same_view_temporal_window=SameViewTemporalWindowConfig(sensor_to_sensor_bounds=(-0.2, 0.0)), + ), + ) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_single_camera_robot_pack_has_no_cross_view_pass() -> None: + plan = _prepare(_robot_pack(num_views=1)) + assert plan.num_gen_tokens == 14 + assert plan.cross_view_empty + assert plan.same_view_offsets is not None + assert plan.same_view_offsets.tolist() == [0, 14, 17] + assert plan.same_view_gather is None + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_robot_samples_keep_independent_clocks_and_action_blocks() -> None: + plan = _prepare(_robot_pack(frames=(3, 4))) + assert plan.num_gen_tokens == 95 + assert plan.items_per_sample == (6, 6) + edges = _cross_edges(plan) + assert all((query < 39) == (key < 39) for query, key in edges) + assert (2, 23) in edges and (41, 71) in edges diff --git a/cosmos_framework/model/generator/mot/multiview_attention.md b/cosmos_framework/model/generator/mot/multiview_attention.md index 530203c7f..5840f7063 100644 --- a/cosmos_framework/model/generator/mot/multiview_attention.md +++ b/cosmos_framework/model/generator/mot/multiview_attention.md @@ -133,6 +133,27 @@ mixed-rate geometry raises instead of treating seconds as frames. FlexAttention does not implement these same-view windows. A control-to-sensor bound does not re-enable edges disabled by `control_attends_sensor=False`. +## Separate-camera robot action conditioning + +`multiview_action_conditioning` configures packing for the bimanual RGB/action +layout, including camera codes in padded action channels. The packer records +`PackedSequence.multiview_action_temporal_factor`, so attention identifies this +layout from the pack rather than interpreting every multi-item sample as cameras. + +The shared multiview item builder describes each ragged RGB clip as a separate +camera target and each camera's replicated action rows as `(T-1, 1, factor)`. +An action item shares its camera's view identity and starts at latent frame one; +all full-rate action rows are retained. RGB and action items have sensor roles +and read the shared caption, even though the actions are clean conditioning. + +These items use the same maskless plan builder and `multiview_attention` settings +as camera-major RGB, LiDAR/radar, and camera-pose control layouts. Cross-view +temporal windows and first-frame inclusion therefore share one configuration +path, with same-view keys excluded from cross-view passes. The separate-camera action layout still requires maskless +attention and baseline geometry: geometry encodings, cross-view bands, and +same-view neighborhood windows are explicitly rejected. Ordinary single-camera +packs without layout metadata continue to use two-way attention. + ## Generation With Camera-Trajectory Controls MultiCamVideo forward-dynamics generation packs a target vision item and one diff --git a/cosmos_framework/model/generator/mot/multiview_attention.py b/cosmos_framework/model/generator/mot/multiview_attention.py index 518227756..c3b915239 100644 --- a/cosmos_framework/model/generator/mot/multiview_attention.py +++ b/cosmos_framework/model/generator/mot/multiview_attention.py @@ -314,6 +314,9 @@ def resolve_multiview_backend( "cross_view_temporal_window.includes_first_frame requires backend='maskless'; the flex " "mask has no first-frame rule" ) + for flag in ("maskless_cross_view_attends_action", "maskless_cross_view_action_attends_action"): + if getattr(config, flag, False): + raise ValueError(f"{flag} requires backend='maskless'; the flex mask has no cross-view action rule") configured_same_view_windows = { name: getattr(config.mask.same_view_temporal_window, name) diff --git a/cosmos_framework/model/generator/mot/multiview_maskless_attention.py b/cosmos_framework/model/generator/mot/multiview_maskless_attention.py index d8976b2ce..2dc676a4d 100644 --- a/cosmos_framework/model/generator/mot/multiview_maskless_attention.py +++ b/cosmos_framework/model/generator/mot/multiview_maskless_attention.py @@ -44,12 +44,13 @@ apply_prope_pairs, prope_pair_channels, ) + +from cosmos_framework.model.generator.mot.merge_attention import merge_attentions_ac_safe from cosmos_framework.model.generator.mot.merge_bridge import ( BridgeFn, DisjointQueriesBridge, MergeAttentionsBridge, ) -from cosmos_framework.model.generator.mot.merge_attention import merge_attentions_ac_safe from cosmos_framework.model.generator.mot.rigrope import apply_mrope_rotary, apply_rig_rotary from cosmos_framework.data.generator.sequence_packing.runtime import ( SequencePack, @@ -157,6 +158,42 @@ def maskless_unavailable_reason(config: MultiviewAttentionConfig) -> str | None: "over each whole sample rather than a partition of one, so it is a mask rule here " "and not a fold" ) + same_view_windows = sorted( + name + for name in ( + "sensor_to_sensor_bounds", + "sensor_to_control_bounds", + "control_to_control_bounds", + "control_to_sensor_bounds", + ) + if getattr(config.mask.same_view_temporal_window, name) is not None + ) + for flag in ("maskless_cross_view_attends_action", "maskless_cross_view_action_attends_action"): + if not getattr(config, flag, False): + continue + if config.mask.attention_scope != "decomposed": + return ( + f"{flag} adds action tokens to the capture-time cross-view rectangles, so it needs " + f"attention_scope='decomposed'; this config has attention_scope={config.mask.attention_scope!r}" + ) + if ( + config.maskless_cross_view_band_radius + or config.maskless_cross_view_include_frame_zero + or config.geometry_position_encoding != "baseline" + ): + return ( + f"{flag} supports neither cross-view band expansion nor RigRoPE/PRoPE geometry, which " + "have no rule for action tokens; keep maskless_cross_view_band_radius=0, " + "maskless_cross_view_include_frame_zero=False and geometry_position_encoding='baseline'" + ) + if same_view_windows: + # Same-view windows need every item on frame offset zero, and the flag times a + # frame-rate action from video frame 1, so every action batch it acts on would be + # refused mid-run. + return ( + f"{flag} acts on camera-pose action batches, which same-view temporal windows do " + f"not serve; this config sets same_view_temporal_window.{same_view_windows}" + ) return None @@ -402,6 +439,8 @@ def _cross_view_partitions( cross_view_temporal_window: CrossViewTemporalWindowConfig, *, include_frame_zero: bool = False, + key_only_cells: dict[tuple[int, float], dict[int, list[tuple[int, int]]]] | None = None, + control_query_cells: dict[tuple[int, float], dict[int, list[tuple[int, int]]]] | None = None, ) -> tuple[CrossViewPartition, ...]: """Batch capture-time rectangles that exclude same-view keys. @@ -428,39 +467,84 @@ def _cross_view_partitions( for every query time on top of the window. It edits that same boolean table rather than adding a rectangle, so a first frame already inside the window is admitted once, and the view split still keeps own-view keys out. + + ``key_only_cells`` has the same layout and adds key spans that never query, e.g. action + tokens on their camera's view group. They join the key side of every rectangle their + instant or capture time reaches, so the view split keeps a query's own view out of them too. + + ``control_query_cells`` holds control spans that query their own rectangles, keyed only by + the other views' spans in that same table under the same instant or window rule. They are + separate groups from the sensor ones, so a token still queries at most once per level. """ # Each view has separate query/key spans: identical for an instant, different # for a directed window. Missing views on either side have empty spans. groups: list[list[tuple[list[tuple[int, int]], list[tuple[int, int]]]]] = [] bounds = temporal_window_bounds(cross_view_temporal_window.bounds) + key_only = key_only_cells or {} if bounds is None: - groups = [[(spans, spans) for spans in views.values()] for views in instant_cells.values() if len(views) > 1] + for instant, queries in instant_cells.items(): + instant_keys = {view: list(spans) for view, spans in queries.items()} + for view, spans in key_only.get(instant, {}).items(): + instant_keys.setdefault(view, []).extend(spans) + views = sorted(instant_keys) + if len(views) > 1: + groups.append([(queries.get(view, []), instant_keys[view]) for view in views]) + groups.extend( + [ + [(spans, spans) for spans in views.values()] + for views in (control_query_cells or {}).values() + if len(views) > 1 + ] + ) else: - samples: dict[int, dict[float, dict[int, list[tuple[int, int]]]]] = {} - for (sample, timestamp), views in instant_cells.items(): - samples.setdefault(sample, {})[timestamp] = views start, end = torch.tensor(bounds, dtype=torch.float32, device="cpu").unbind() # each [] - for cells in samples.values(): - timestamps = sorted(cells) - times = torch.tensor(timestamps, dtype=torch.float32, device="cpu") # [F_union] - offsets = times[None, :] - times[:, None] # [F_union,F_union], key time minus query time - allowed = (offsets >= start) & (offsets <= end) # [F_union,F_union] - if cross_view_temporal_window.includes_first_frame: - # Column 0 is the sample's earliest capture time across its streams; - # without frame offsets, every stream starts at capture time 0. - allowed[:, 0] = True # [F_union] - if include_frame_zero and 0 in cells: - # Experimental bands refer to instant zero, even for offset clips. - allowed[:, timestamps.index(0)] = True # [F_union] - for row, timestamp in enumerate(timestamps): - queries = cells[timestamp] - keys: dict[int, list[tuple[int, int]]] = {} - for column in allowed[row].nonzero(as_tuple=True)[0].tolist(): - for view, spans in cells[timestamps[column]].items(): - keys.setdefault(view, []).extend(spans) - views = sorted(queries.keys() | keys.keys()) - if len(views) > 1: - groups.append([(queries.get(view, []), keys.get(view, [])) for view in views]) + + def _by_sample( + cells: dict[tuple[int, float], dict[int, list[tuple[int, int]]]], + ) -> dict[int, dict[float, dict[int, list[tuple[int, int]]]]]: + samples: dict[int, dict[float, dict[int, list[tuple[int, int]]]]] = {} + for (sample, timestamp), views in cells.items(): + samples.setdefault(sample, {})[timestamp] = views + return samples + + sensor_samples = _by_sample(instant_cells) + key_only_samples = _by_sample(key_only) + control_samples = _by_sample(control_query_cells or {}) + # Sensor queries key sensor and key-only spans; control queries key their own table. + for query_samples, key_sources, first_frame_column in ( + (sensor_samples, (sensor_samples, key_only_samples), True), + (control_samples, (control_samples,), False), + ): + for sample, cells in query_samples.items(): + sample_keys = [source.get(sample, {}) for source in key_sources] + timestamps = sorted(cells) + key_timestamps = sorted(set().union(*(source.keys() for source in sample_keys))) + times = torch.tensor(timestamps, dtype=torch.float32, device="cpu") # [F_query] + key_times = torch.tensor(key_timestamps, dtype=torch.float32, device="cpu") # [F_key] + offsets = key_times[None, :] - times[:, None] # [F_query,F_key], key time minus query time + allowed = (offsets >= start) & (offsets <= end) # [F_query,F_key] + if cross_view_temporal_window.includes_first_frame: + if first_frame_column: + # Column 0 is the sample's earliest capture time across its streams; + # without frame offsets, every stream starts at capture time 0. + allowed[:, 0] = True # [F_key] + else: + # The sample's first frame is capture time 0; a control stream that starts + # later, like actions from video frame 1, has no row there to admit. + allowed |= (key_times == 0.0)[None, :] # [F_query,F_key] + if include_frame_zero and 0 in key_timestamps: + # Experimental bands refer to instant zero, even for offset clips. + allowed[:, key_timestamps.index(0)] = True # [F_key] + for row, timestamp in enumerate(timestamps): + queries = cells[timestamp] + keys: dict[int, list[tuple[int, int]]] = {} + for column in allowed[row].nonzero(as_tuple=True)[0].tolist(): + for source in sample_keys: + for view, spans in source.get(key_timestamps[column], {}).items(): + keys.setdefault(view, []).extend(spans) + views = sorted(queries.keys() | keys.keys()) + if len(views) > 1: + groups.append([(queries.get(view, []), keys.get(view, [])) for view in views]) partitions: list[CrossViewPartition] = [] while groups: @@ -722,6 +806,8 @@ def build_multiview_maskless_plan( cross_view_band_radius: int = 0, include_frame_zero: bool = False, cross_view_band_keys: str = "other_views", + cross_view_key_only: Sequence[bool] | None = None, + cross_view_control_queries: Sequence[bool] | None = None, ) -> MultiviewMasklessPlan: """Describe a batch's GEN stream to :func:`multiview_maskless_attention`. @@ -745,7 +831,10 @@ def build_multiview_maskless_plan( partitioned into rectangles between distinct views. Every control rule in the mask is a view rule, never an instant one, so a control token is absent from this partition as a key and as a query alike. It is sliced out rather than - masked, so nothing is computed for it and nothing reaches it. + masked, so nothing is computed for it and nothing reaches it. The exceptions are a + ``cross_view_key_only`` item, which joins the sensor rectangles as keys but never queries + them, and a ``cross_view_control_queries`` item, which queries rectangles of its own keyed + by the other views' tokens of such items. A sample owning a *single* same-view group is sliced out of the cross-instant partition on the same terms, and for a sharper reason: its instant groups are then subsets of its one @@ -824,7 +913,16 @@ def build_multiview_maskless_plan( keys already covered by the same-view pass are not counted twice. With no expansion, the default is accepted and all cross-view rectangles still exclude own-view keys. - + cross_view_key_only: per item, whether a control item also serves as a cross-view + *key*, never a query: a sensor query then reaches it on the other views at the + instants or capture times its window admits, the way it reaches their sensor + tokens. Its frame ``f`` is timed ``(f + frame_offset) * seconds_per_frame``, as a + sensor item's is. Needs ``attention_scope="decomposed"`` and no band expansion. + ``None`` keeps every control item out of that pass. + cross_view_control_queries: per item, whether a control item's tokens also *query* the + cross-view pass, keyed only by the other views' tokens of the items so marked, under + the same instant or window rule and the same clock as ``cross_view_key_only``. They + never key sensor tokens there. Same requirements as ``cross_view_key_only``. Returns: The plan, with its partitions as varlen offsets and the gathers that reach them. @@ -897,6 +995,8 @@ def build_multiview_maskless_plan( rates = [1.0] * num_items if seconds_per_frame is None else list(seconds_per_frame) offsets = [0] * num_items if frame_offsets is None else list(frame_offsets) control = [False] * num_items if is_control is None else list(is_control) + key_only = [False] * num_items if cross_view_key_only is None else list(cross_view_key_only) + control_queries = [False] * num_items if cross_view_control_queries is None else list(cross_view_control_queries) axes = [0] * num_items if view_axis is None else list(view_axis) # ``None`` reproduces what the axis used to decide on its own: the cameras' axis is the one # captions are written for, and every other axis is a sensor that reads all of them. Callers @@ -909,11 +1009,23 @@ def build_multiview_maskless_plan( ("seconds_per_frame", rates), ("frame_offsets", offsets), ("is_control", control), + ("cross_view_key_only", key_only), + ("cross_view_control_queries", control_queries), ("view_axis", axes), ("caption_access", accesses), ): if len(values) != num_items: raise ValueError(f"num_views describes {num_items} items but {name} describes {len(values)}.") + for name, marks in (("cross_view_key_only", key_only), ("cross_view_control_queries", control_queries)): + if not any(marks): + continue + if any(mark and not is_ctrl for mark, is_ctrl in zip(marks, control)): + raise ValueError(f"{name} marks a sensor item, which already queries and keys that pass.") + if attention_scope != "decomposed" or expands_band: + raise ValueError( + f"{name} adds control tokens to the cross-view rectangles, so it needs " + "attention_scope='decomposed' and no cross-view band expansion." + ) if any(rate <= 0 or not math.isfinite(rate) for rate in rates): raise ValueError(f"seconds_per_frame must be positive and finite, got {rates}.") same_view = same_view_temporal_window or SameViewTemporalWindowConfig() @@ -1009,6 +1121,10 @@ def build_multiview_maskless_plan( group_access: dict[int, CaptionAccess] = {} view_spans: list[list[tuple[int, int]]] = [] instant_cells: dict[tuple[int, float], dict[int, list[tuple[int, int]]]] = {} + # The same layout for ``cross_view_key_only`` control items, which key those rectangles only. + key_only_cells: dict[tuple[int, float], dict[int, list[tuple[int, int]]]] = {} + # And for ``cross_view_control_queries`` items, which query and key their own rectangles. + control_query_cells: dict[tuple[int, float], dict[int, list[tuple[int, int]]]] = {} caption_reader_runs: list[torch.Tensor] = [] caption_reader_lens: list[int] = [] # Each same-view group's tokens kept apart by what they are, for the split @@ -1085,7 +1201,24 @@ def build_multiview_maskless_plan( f"within each view group; group={key}, periods={period!r} and {rates[item]!r}." ) - if not control[item] and not single_group_sample[sample] and attention_scope != "same_view": + if ( + (not control[item] or key_only[item] or control_queries[item]) + and not single_group_sample[sample] + and attention_scope != "same_view" + ): + # A marked control item writes to its own tables, never to the sensor queries'. + destinations = ( + [ + cells + for cells, marked in ( + (key_only_cells, key_only[item]), + (control_query_cells, control_queries[item]), + ) + if marked + ] + if control[item] + else [instant_cells] + ) # The epsilon nudges a frame whose midpoint lands exactly on an anchor boundary # into the later group rather than leaving it to float rounding. Exact landings # need commensurate spans; at 10Hz against 7.5Hz the margin is an eighth of a @@ -1106,9 +1239,10 @@ def build_multiview_maskless_plan( group = view_group[(sample, axes[item], view)] for frame, instant in enumerate(instants): start = position + (view * frames + frame) * spatial - instant_cells.setdefault((sample, instant), {}).setdefault(group, []).append( - (start, start + spatial) - ) + for cells in destinations: + cells.setdefault((sample, instant), {}).setdefault(group, []).append( + (start, start + spatial) + ) # The cameras' axis is the one captions are written for; every other axis is a # sensor no caption describes, which is what the flag decides the fate of. if subset_gen_to_und and accesses[item] != "no_captions": @@ -1187,7 +1321,14 @@ def _gen_to_und_subset() -> tuple[torch.Tensor | None, int]: if expands_band else cross_view ) - partitions = _cross_view_partitions(instant_cells, device, partition_window, include_frame_zero=include_frame_zero) + partitions = _cross_view_partitions( + instant_cells, + device, + partition_window, + include_frame_zero=include_frame_zero, + key_only_cells=key_only_cells, + control_query_cells=control_query_cells, + ) caption_gather, caption_lens = _caption_partition(captions, view_group, group_access, device, bool(pad_tokens)) caption_q_gather, caption_q_offsets, caption_q_max_len, caption_offsets, caption_max_len = ( (None, None, 0, None, 0) diff --git a/cosmos_framework/model/generator/mot/multiview_maskless_exact_test.py b/cosmos_framework/model/generator/mot/multiview_maskless_exact_test.py index e3779d5ae..fd90a5157 100644 --- a/cosmos_framework/model/generator/mot/multiview_maskless_exact_test.py +++ b/cosmos_framework/model/generator/mot/multiview_maskless_exact_test.py @@ -149,25 +149,34 @@ def _batch( per_view_captions: bool, window: TemporalWindow | None = None, include_first_frame: bool = False, + control_keys: bool = False, + control_queries: bool = False, ) -> tuple[list[torch.Tensor], list[SequencePack], maskless.MultiviewMasklessPlan, list[Token], list[tuple[int, int]]]: - # Item = (views, frames, spatial tokens, modality axis, control, seconds/frame). - rgb = (7, 5, 2, 0, False, 2 / 15) + # Item = (views, frames, spatial tokens, modality axis, control, seconds/frame, frame offset). + rgb = (7, 5, 2, 0, False, 2 / 15, 0) if case == "rgb": samples = [[rgb]] elif case == "rgb11": - samples = [[(11, 5, 2, 0, False, 2 / 15)]] + samples = [[(11, 5, 2, 0, False, 2 / 15, 0)]] elif case == "single": - samples = [[(1, 3, 2, 0, False, 2 / 15)]] + samples = [[(1, 3, 2, 0, False, 2 / 15, 0)]] elif case == "ragged_joint": samples = [ [ - (3, 3, 2, 0, True, 2 / 15), - (3, 3, 2, 0, False, 2 / 15), - (1, 4, 1, 1, True, 0.1), - (1, 4, 1, 1, False, 0.1), + (3, 3, 2, 0, True, 2 / 15, 0), + (3, 3, 2, 0, False, 2 / 15, 0), + (1, 4, 1, 1, True, 0.1, 0), + (1, 4, 1, 1, False, 0.1, 0), ], - [(1, 2, 3, 0, False, 2 / 15)], - [(2, 4, 1, 0, False, 2 / 15)], + [(1, 2, 3, 0, False, 2 / 15, 0)], + [(2, 4, 1, 0, False, 2 / 15, 0)], + ] + elif case == "rgb_action": + # A camera-pose action control after its camera target: frame-rate rows from video + # frame 1 on, against latent frames four video frames apart, and a single-camera sample. + samples = [ + [(3, 3, 2, 0, False, 4 / 30, 0), (3, 8, 1, 0, True, 1 / 30, 1)], + [(1, 3, 2, 0, False, 4 / 30, 0), (1, 8, 1, 0, True, 1 / 30, 1)], ] else: raise ValueError(case) @@ -184,20 +193,22 @@ def _batch( captions.append(caption) text_layout.extend((sample, view) for view, length in caption for _ in range(length)) und_len = sum(length for _, length in caption) - gen_len = sum(views * frames * spatial for views, frames, spatial, _, _, _ in items) + gen_len = sum(views * frames * spatial for views, frames, spatial, *_ in items) split_lens.extend((und_len, gen_len)) sample_lens.append(und_len + gen_len) und_indices.extend(range(cursor, cursor + und_len)) gen_indices.extend(range(cursor + und_len, cursor + und_len + gen_len)) cursor += und_len + gen_len - for views, frames, spatial, axis, control, period in items: - frame_times = (torch.arange(frames, dtype=torch.float32, device="cpu") * period).tolist() # [F] -> list + for views, frames, spatial, axis, control, period, offset in items: + frame_times = ( + (torch.arange(frames, dtype=torch.float32, device="cpu") + offset) * period + ).tolist() # [F] -> list for view in range(views): for frame in range(frames): instant = ( frame_times[frame] if window is not None - else math.floor((frame + 0.5) * period / items[0][5] + 1e-6) + else math.floor((frame + offset + 0.5) * period / items[0][5] + 1e-6) ) layout.extend([(sample, axis, view, instant, control)] * spatial) tensors = [ @@ -221,13 +232,16 @@ def _batch( items = [item for sample_items in samples for item in sample_items] plan = maskless.build_multiview_maskless_plan( [item[0] for item in items], - [(views * frames, 1, spatial) for views, frames, spatial, _, _, _ in items], + [(views * frames, 1, spatial) for views, frames, spatial, *_ in items], device=device, items_per_sample=[len(sample_items) for sample_items in samples], is_control=[item[4] for item in items], control_attends_sensor=control_attends_sensor, view_axis=[item[3] for item in items], seconds_per_frame=[item[5] for item in items], + frame_offsets=[item[6] for item in items], + cross_view_key_only=[item[4] for item in items] if control_keys else None, + cross_view_control_queries=[item[4] for item in items] if control_queries else None, captions=captions if per_view_captions else None, padded_gen_tokens=get_full_only_seq(packs[0])[0].shape[0], cross_view_temporal_window=CrossViewTemporalWindowConfig( @@ -243,6 +257,8 @@ def _pair_multiplicity( control_attends_sensor: bool, window: TemporalWindow | None = None, include_first_frame: bool = False, + control_keys: bool = False, + control_queries: bool = False, ) -> int: qs, qa, qv, qt, qc = query ks, ka, kv, kt, kc = key @@ -254,7 +270,10 @@ def _pair_multiplicity( in_time = window[0] - 1e-4 <= kt - qt <= window[1] + 1e-4 or (include_first_frame and kt == 0.0) if qs != ks: return 0 - count = int(same_view and (not qc or kc or control_attends_sensor)) + int(not qc and not kc and in_time) + # Control tokens query the cross-view pass only with ``control_queries``, and then key only + # other control tokens; with ``control_keys`` they key sensor queries. + cross_pair = (not qc and (not kc or control_keys)) or (qc and kc and control_queries) + count = int(same_view and (not qc or kc or control_attends_sensor)) + int(cross_pair and in_time) return int(count > 0) @@ -265,6 +284,8 @@ def _dense_reference( control_attends_sensor: bool, window: TemporalWindow | None = None, include_first_frame: bool = False, + control_keys: bool = False, + control_queries: bool = False, ) -> torch.Tensor: # [N_gen,H*D] q, k, v = [get_full_only_seq(pack)[0][: len(layout)].double() for pack in packs] # each [N_gen,H,D] text_k, text_v = [get_causal_seq(pack)[0][: len(text_layout)].double() for pack in packs[1:]] # each [N_text,Hkv,D] @@ -279,6 +300,8 @@ def _dense_reference( control_attends_sensor, window, include_first_frame, + control_keys, + control_queries, ) for key in layout ] @@ -312,6 +335,61 @@ def test_maskless_forward_and_backward( ) -> None: if include_first_frame and window is None: pytest.skip("The first-frame rule widens a window; test_first_frame_rule_requires_a_window covers None.") + _check_forward_backward_and_edges( + monkeypatch, + case, + control_attends_sensor, + per_view_captions, + window, + include_first_frame, + ) + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("control_keys", [True, False]) +@pytest.mark.parametrize("control_queries", [True, False]) +@pytest.mark.parametrize("control_attends_sensor", [True, False]) +@pytest.mark.parametrize("per_view_captions", [True, False]) +@pytest.mark.parametrize( + "window", [None, (0.0, 0.0), (-0.04, 0.04), (-0.1, 0.0), (-0.67, 0.67), (0.05, 0.2), (10.0, 11.0)] +) +@pytest.mark.parametrize("include_first_frame", [False, True]) +def test_cross_view_action_keys_forward_and_backward( + monkeypatch: pytest.MonkeyPatch, + control_keys: bool, + control_queries: bool, + control_attends_sensor: bool, + per_view_captions: bool, + window: TemporalWindow | None, + include_first_frame: bool, +) -> None: + """Action rows reach other views by capture time: as camera queries' keys, and as queries of each other.""" + if window is None and include_first_frame: + pytest.skip("The first-frame rule widens a window; test_first_frame_rule_requires_a_window covers None.") + _check_forward_backward_and_edges( + monkeypatch, + "rgb_action", + control_attends_sensor, + per_view_captions, + window, + include_first_frame, + control_keys=control_keys, + control_queries=control_queries, + ) + + +def _check_forward_backward_and_edges( + monkeypatch: pytest.MonkeyPatch, + case: str, + control_attends_sensor: bool, + per_view_captions: bool, + window: TemporalWindow | None, + include_first_frame: bool, + *, + control_keys: bool = False, + control_queries: bool = False, +) -> None: monkeypatch.setattr(maskless, "attention", _reference_attention) monkeypatch.setattr(maskless, "merge_attentions", _reference_merge) torch.manual_seed(51) @@ -323,10 +401,12 @@ def test_maskless_forward_and_backward( per_view_captions=per_view_captions, window=window, include_first_frame=include_first_frame, + control_keys=control_keys, + control_queries=control_queries, ) actual = maskless.multiview_maskless_gen_attention(*packs, plan=plan) # [N_padded,H*D] expected = _dense_reference( - packs, layout, text_layout, control_attends_sensor, window, include_first_frame + packs, layout, text_layout, control_attends_sensor, window, include_first_frame, control_keys, control_queries ) # [N_gen,H*D] torch.testing.assert_close(actual[: len(layout)], expected, atol=1e-10, rtol=1e-10) assert torch.count_nonzero(actual[len(layout) :]) == 0 @@ -346,7 +426,9 @@ def test_maskless_forward_and_backward( edges: set[tuple[int, int]] = set() for partition in plan.cross_view_partitions: assert len(set(partition.query_indices.tolist())) == partition.query_indices.numel() - if window is None: + if window is None and not (control_keys and control_queries): + # Instants are disjoint, so a key repeats only when control rows key both an + # instant's camera queries and its control queries. assert len(set(partition.key_indices.tolist())) == partition.key_indices.numel() for qs, qe, ks, ke in zip( partition.query_offsets[:-1], @@ -363,12 +445,13 @@ def test_maskless_forward_and_backward( (qi, ki) for qi, query in enumerate(layout) for ki, key in enumerate(layout) - if not query[4] - and not key[4] + if ((not query[4] and (not key[4] or control_keys)) or (query[4] and key[4] and control_queries)) and group_counts[query[0]] > 1 and query[1:3] != key[1:3] # Exclude the same-view contribution to isolate this temporal pass. - and _pair_multiplicity(query, key, control_attends_sensor, window, include_first_frame) + and _pair_multiplicity( + query, key, control_attends_sensor, window, include_first_frame, control_keys, control_queries + ) } assert edges == expected_edges max_depth = math.ceil(math.log2(max(group_counts.values()))) @@ -432,6 +515,102 @@ def test_maskless_compiles_as_one_graph( torch.testing.assert_close(actual, expected, atol=0, rtol=0) +@pytest.mark.L0 +@pytest.mark.CPU +def test_cross_view_action_flags_are_opt_in() -> None: + config = MultiviewAttentionConfig(backend="maskless") + assert config.maskless_cross_view_attends_action is False + assert config.maskless_cross_view_action_attends_action is False + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("name", ["cross_view_key_only", "cross_view_control_queries"]) +def test_cross_view_control_items_are_validated(name: str) -> None: + plan_args: dict[str, Any] = dict( + num_views=[2, 2], + token_shapes=[(4, 1, 1), (6, 1, 1)], + device=torch.device("cpu"), + seconds_per_frame=[0.1, 0.05], + items_per_sample=[2], + is_control=[False, True], + control_attends_sensor=True, + ) + with pytest.raises(ValueError, match="marks a sensor item"): + maskless.build_multiview_maskless_plan(**plan_args, **{name: [True, True]}) + for refused in ({"attention_scope": "same_view"}, {"cross_view_band_radius": 1}, {"include_frame_zero": True}): + with pytest.raises(ValueError, match=f"{name} adds control tokens"): + maskless.build_multiview_maskless_plan(**plan_args, **refused, **{name: [False, True]}) + with pytest.raises(ValueError, match="describes 1"): + maskless.build_multiview_maskless_plan(**plan_args, **{name: [True]}) + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("flag", ["maskless_cross_view_attends_action", "maskless_cross_view_action_attends_action"]) +def test_cross_view_action_flags_are_refused_where_unsupported(flag: str) -> None: + for window in (None, CrossViewTemporalWindowConfig(bounds=(-0.67, 0.67), includes_first_frame=True)): + supported = MultiviewAttentionConfig( + backend="maskless", + mask=MultiviewAttentionMaskConfig( + attention_scope="decomposed", cross_view_temporal_window=window or CrossViewTemporalWindowConfig() + ), + **{flag: True}, + ) + assert maskless.maskless_unavailable_reason(supported) is None + unsupported = [ + MultiviewAttentionConfig( + backend="maskless", + mask=MultiviewAttentionMaskConfig(attention_scope="same_view"), + **{flag: True}, + ), + MultiviewAttentionConfig( + backend="maskless", + geometry_position_encoding="rigrope_cross_view", + mask=MultiviewAttentionMaskConfig(attention_scope="decomposed"), + **{flag: True}, + ), + MultiviewAttentionConfig( + backend="maskless", + maskless_cross_view_band_radius=1, + mask=MultiviewAttentionMaskConfig(attention_scope="decomposed"), + **{flag: True}, + ), + MultiviewAttentionConfig( + backend="maskless", + maskless_cross_view_include_frame_zero=True, + mask=MultiviewAttentionMaskConfig(attention_scope="decomposed"), + **{flag: True}, + ), + *( + MultiviewAttentionConfig( + backend="maskless", + mask=MultiviewAttentionMaskConfig( + attention_scope="decomposed", + same_view_temporal_window=SameViewTemporalWindowConfig(**{bounds: (-0.4, 0.4)}), + ), + **{flag: True}, + ) + for bounds in ( + "sensor_to_sensor_bounds", + "sensor_to_control_bounds", + "control_to_control_bounds", + "control_to_sensor_bounds", + ) + ), + ] + for config in unsupported: + reason = maskless.maskless_unavailable_reason(config) + assert reason is not None and flag in reason + for preference in ("auto", "flex_triton", "flex_flash"): + with pytest.raises(ValueError, match=f"{flag} requires backend='maskless'"): + resolve_multiview_backend( + torch.device("cpu"), + preference, + config=MultiviewAttentionConfig(backend=preference, **{flag: True}), + ) + + @pytest.mark.L0 @pytest.mark.CPU def test_first_frame_rule_requires_a_window() -> None: diff --git a/cosmos_framework/model/generator/mot/parallelize_unified_mot.py b/cosmos_framework/model/generator/mot/parallelize_unified_mot.py index e057da118..868929068 100644 --- a/cosmos_framework/model/generator/mot/parallelize_unified_mot.py +++ b/cosmos_framework/model/generator/mot/parallelize_unified_mot.py @@ -34,12 +34,12 @@ from cosmos_framework.configs.base.defaults.activation_checkpointing import ActivationCheckpointingConfig from cosmos_framework.configs.base.defaults.compile import CompileConfig -from cosmos_framework.data.generator.sequence_packing.runtime import SequencePack from cosmos_framework.model.generator.mot.activation_marks import enable_marking, is_mark_op from cosmos_framework.model.generator.mot.attention import SplitInfo, dispatch_attention from cosmos_framework.model.generator.mot.context_parallel_utils import context_parallel_attention from cosmos_framework.model.generator.mot.replicated_io import apply_replicated_attention_io_cp from cosmos_framework.model.generator.utils.memory import KVToStore, MemoryValue +from cosmos_framework.data.generator.sequence_packing.runtime import SequencePack from cosmos_framework.utils.generator.activation_offloading import ( checkpoint_with_flattened_inputs, offload_checkpoint_inputs, @@ -512,6 +512,12 @@ def apply_cp( cp_mesh = parallel_dims.cp_mesh for _, block in model.model.layers.named_children(): attn = block.self_attn + # Ulysses scatters query heads across ranks; KV heads can be repeated for GQA. + if attn.num_attention_heads % parallel_dims.cp != 0: + raise ValueError( + f"Query heads ({attn.num_attention_heads}) must be divisible by " + f"CP (Context Parallelism) size ({parallel_dims.cp})." + ) attn.dispatch_attention_fn = ContextParallelDispatch( cp_mesh, wrapped_dispatch=attn.dispatch_attention_fn, diff --git a/cosmos_framework/model/generator/mot/parallelize_unified_mot_test.py b/cosmos_framework/model/generator/mot/parallelize_unified_mot_test.py index 300408b7e..9b4e2086f 100644 --- a/cosmos_framework/model/generator/mot/parallelize_unified_mot_test.py +++ b/cosmos_framework/model/generator/mot/parallelize_unified_mot_test.py @@ -40,13 +40,16 @@ from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.model.generator.mot.parallelize_unified_mot import ( _PACK_DYNAMIC_LEN_KEYS, + ContextParallelDispatch, _mark_pack_unbacked, _wrap_forward_with_unbacked_pack, apply_compile, + apply_cp, apply_fsdp, materialize_non_offloaded_state, parallelize_unified_mot, ) +from cosmos_framework.utils.generator.parallelism import ParallelDims pytestmark = [pytest.mark.L0, pytest.mark.CPU] @@ -272,6 +275,43 @@ def test_defaults_to_not_marking(self) -> None: assert CompileConfig().mark_unbacked is False +class TestContextParallelHeads: + @staticmethod + def _model_with_query_heads(num_heads: int) -> torch.nn.Module: + attention = torch.nn.Module() + attention.num_attention_heads = num_heads + attention.dispatch_attention_fn = torch.nn.Identity() + block = torch.nn.Module() + block.self_attn = attention + inner = torch.nn.Module() + inner.layers = torch.nn.ModuleList([block]) + model = torch.nn.Module() + model.model = inner + return model + + @pytest.mark.parametrize("cp", [32, 64]) + def test_accepts_cp_dividing_query_heads(self, cp: int) -> None: + model = self._model_with_query_heads(64) + parallel_dims = ParallelDims(world_size=128, dp_shard=128, cp=cp) + parallel_dims._meshes["cp"] = object() + original_dispatch = model.model.layers[0].self_attn.dispatch_attention_fn + + assert apply_cp(model, parallel_dims) is model + + dispatch = model.model.layers[0].self_attn.dispatch_attention_fn + assert isinstance(dispatch, ContextParallelDispatch) + assert dispatch.cp_mesh is parallel_dims.cp_mesh + assert dispatch.wrapped_dispatch is original_dispatch + + @pytest.mark.parametrize(("num_heads", "cp"), [(32, 64), (48, 32)]) + def test_rejects_cp_not_dividing_query_heads(self, num_heads: int, cp: int) -> None: + model = self._model_with_query_heads(num_heads) + parallel_dims = ParallelDims(world_size=128, dp_shard=128, cp=cp) + + with pytest.raises(ValueError, match="Query heads .* must be divisible by CP"): + apply_cp(model, parallel_dims) + + class TestFSDPCPUOffload: @staticmethod def _model_with_two_layers() -> torch.nn.Module: diff --git a/cosmos_framework/model/generator/mot/three_way_attention_test.py b/cosmos_framework/model/generator/mot/three_way_attention_test.py index 53480c3bd..d72d3c68e 100644 --- a/cosmos_framework/model/generator/mot/three_way_attention_test.py +++ b/cosmos_framework/model/generator/mot/three_way_attention_test.py @@ -1179,17 +1179,11 @@ def test_attention_empty_varlen_self_attention_returns_zero(dtype: torch.dtype): ``[0, 0]``) but padded to a fixed buffer size should return a zero output, not NaN. - Mirrors the text self-attention call inside - ``three_way_attention_with_memory`` when no caption is present: - Q / K / V come from the pack's ``causal_seq`` slot (padded to a - constant length for compile stability), and ``text_kv_offsets`` is - ``[0, 0]`` (or, if ``clamp_empty_varlen_kv`` is True, ``[0, 1]`` — - this test pins the unclamped case). The output is passed straight - to ``from_mode_splits``; positions beyond the real range are - discarded, but they must be *zero*, not NaN, to avoid contaminating - downstream tensors. The LSE is not requested here — the text SA - output does not feed ``merge_attentions``, so its LSE value is - irrelevant. + This checks the raw attention kernel's empty-range contract. + Production text self-attention in ``three_way_attention_with_memory`` + instead evaluates padded queries against at least one KV row, then + masks padded outputs to avoid uninitialized backward workspace. + The LSE is not requested here; this test only checks the output. """ device = torch.device("cuda") diff --git a/cosmos_framework/model/generator/mot/unified_mot.py b/cosmos_framework/model/generator/mot/unified_mot.py index a71f49b6c..aa091e45a 100644 --- a/cosmos_framework/model/generator/mot/unified_mot.py +++ b/cosmos_framework/model/generator/mot/unified_mot.py @@ -13,20 +13,9 @@ from torch import nn from torch.distributed import ProcessGroup -from cosmos_framework.data.generator.sequence_packing.runtime import ( - SequencePack, - from_all_seq, - from_und_gen_splits, - get_gen_seq, - get_num_real_tokens, - get_und_seq, - has_pad_segment, - set_gen_seq, - set_und_seq, - zeros_like, -) from cosmos_framework.model.attention import attention as imaginaire_attention from cosmos_framework.model.attention.masks import CausalType +from cosmos_framework.utils import log from cosmos_framework.model.generator.mot.attention import ( AttentionMaskType, dispatch_attention, @@ -89,7 +78,18 @@ ) from cosmos_framework.model.generator.utils.memory import KVToStore, MemoryState, MemoryValue from cosmos_framework.model.generator.utils.numerics import native_rms_norm -from cosmos_framework.utils import log +from cosmos_framework.data.generator.sequence_packing.runtime import ( + SequencePack, + from_all_seq, + from_und_gen_splits, + get_gen_seq, + get_num_real_tokens, + get_und_seq, + has_pad_segment, + set_gen_seq, + set_und_seq, + zeros_like, +) # Torch optimization settings torch._dynamo.config.cache_size_limit = 512 @@ -746,13 +746,22 @@ def forward( memory_value, ) - q_und_in = self.q_proj(get_und_seq(pack)) # [N_und,num_heads*head_dim] + text_cache = getattr(memory_value, "text_cache", None) + projected_text = text_cache.projected if text_cache is not None else None + cached_text = projected_text is not None + q_und_in = ( + self.q_proj(get_und_seq(pack)) if projected_text is None else projected_text[0].flatten(-2) + ) # [N_und,num_heads*head_dim] q_gen_in = self.q_proj_moe_gen(get_gen_seq(pack)) # [N_gen,num_heads*head_dim] - k_und_in = self.k_proj(get_und_seq(pack)) # [N_und,num_kv_heads*head_dim] + k_und_in = ( + self.k_proj(get_und_seq(pack)) if projected_text is None else projected_text[1].flatten(-2) + ) # [N_und,num_kv_heads*head_dim] k_gen_in = self.k_proj_moe_gen(get_gen_seq(pack)) # [N_gen,num_kv_heads*head_dim] - v_und_in = self.v_proj(get_und_seq(pack)) # [N_und,num_kv_heads*head_dim] + v_und_in = ( + self.v_proj(get_und_seq(pack)) if projected_text is None else projected_text[2].flatten(-2) + ) # [N_und,num_kv_heads*head_dim] v_gen_in = self.v_proj_moe_gen(get_gen_seq(pack)) # [N_gen,num_kv_heads*head_dim] q_und = q_und_in.view(-1, self.num_attention_heads, self.head_dim) # [N_und,num_heads,head_dim] @@ -775,8 +784,9 @@ def forward( for head_split in (q_und, k_und, v_und, q_gen, k_gen, v_gen): torch._dynamo.mark_static(head_split, 1) - q_und = _apply_request_rms_norm(self.q_norm, q_und, memory_value) # [N_und,num_heads,head_dim] - k_und = _apply_request_rms_norm(self.k_norm, k_und, memory_value) # [N_und,num_kv_heads,head_dim] + if not cached_text: + q_und = _apply_request_rms_norm(self.q_norm, q_und, memory_value) # [N_und,num_heads,head_dim] + k_und = _apply_request_rms_norm(self.k_norm, k_und, memory_value) # [N_und,num_kv_heads,head_dim] q_gen = _apply_request_rms_norm(self.q_norm_moe_gen, q_gen, memory_value) # [N_gen,num_heads,head_dim] k_gen = _apply_request_rms_norm(self.k_norm_moe_gen, k_gen, memory_value) # [N_gen,num_kv_heads,head_dim] @@ -784,13 +794,18 @@ def forward( packed_cos = packed_position_embeddings[0] packed_sin = packed_position_embeddings[1] - q_und_, k_und_ = self._apply_rotary_pos_emb( - q_und, - k_und, - get_und_seq(packed_cos), - get_und_seq(packed_sin), - unsqueeze_dim=1, - ) # q_und_: [N_und,num_heads,head_dim], k_und_: [N_und,num_kv_heads,head_dim] + if cached_text: + q_und_, k_und_ = q_und, k_und # each [N_und,H,D], already post-RoPE + else: + q_und_, k_und_ = self._apply_rotary_pos_emb( + q_und, + k_und, + get_und_seq(packed_cos), + get_und_seq(packed_sin), + unsqueeze_dim=1, + ) # q_und_: [N_und,num_heads,head_dim], k_und_: [N_und,num_kv_heads,head_dim] + if text_cache is not None: + text_cache.projected = (q_und_, k_und_, v_und) maskless_plan = getattr(attention_mask, "multiview_maskless", None) defer_gen_rope = maskless_plan is not None and maskless_plan.mrope_cos is not None if defer_gen_rope: @@ -816,14 +831,20 @@ def forward( # qk_norm_for_text=False, i.e. the Nemotron-3 config). The raw k_und_ is kept # in packed_key_states_ for the reasoner's own causal self-attention. if self.k_norm_und_for_gen is not None: - k_und_normalized = self.k_norm_und_for_gen(k_und) # RMSNorm before RoPE - _, k_und_for_gen_ = self._apply_rotary_pos_emb( - q_und_, # dummy q — already RoPE-applied, output discarded; only k is used - k_und_normalized, - get_und_seq(packed_cos), - get_und_seq(packed_sin), - unsqueeze_dim=1, - ) # k_und_for_gen_: [N_und,num_kv_heads,head_dim] + if cached_text: + assert text_cache is not None and text_cache.projected_normalized_k is not None + k_und_for_gen_ = text_cache.projected_normalized_k # [N_und,num_kv_heads,head_dim] + else: + k_und_normalized = self.k_norm_und_for_gen(k_und) # [N_und,num_kv_heads,head_dim], RMSNorm before RoPE + _, k_und_for_gen_ = self._apply_rotary_pos_emb( + q_und_, # dummy q — already RoPE-applied, output discarded; only k is used + k_und_normalized, + get_und_seq(packed_cos), + get_und_seq(packed_sin), + unsqueeze_dim=1, + ) # k_und_for_gen_: [N_und,num_kv_heads,head_dim] + if text_cache is not None: + text_cache.projected_normalized_k = k_und_for_gen_ packed_key_states_normalized_: SequencePack | None = from_und_gen_splits(k_und_for_gen_, k_gen_, pack) else: packed_key_states_normalized_ = None @@ -912,7 +933,13 @@ def forward( cp_group, ) # [N_gen,hidden_size] else: - und_seq = self.o_proj(get_und_seq(packed_attn_output)) # [N_und,hidden_size] + if cached_text: + assert text_cache is not None and text_cache.hidden_output is not None + # The layer restores its cached final caption output; this value + # only keeps the existing packed geometry during media execution. + und_seq = text_cache.hidden_output # [N_und,hidden_size] + else: + und_seq = self.o_proj(get_und_seq(packed_attn_output)) # [N_und,hidden_size] gen_seq = self.o_proj_moe_gen(get_gen_seq(packed_attn_output)) # [N_gen,hidden_size] return from_und_gen_splits(und_seq, gen_seq, pack), kv_to_store # [N_und+N_gen,hidden_size] @@ -1168,6 +1195,10 @@ def _impl_forward( for i, decoder_layer in enumerate(self.layers): # MemoryState: produce read-only MemoryValue for this layer (outside compile) memory_value = memory.read_for_layer(i) if memory is not None else None + text_cache_layers = getattr(attention_mask, "ar_text_cache_layers", None) + if text_cache_layers is not None: + assert memory_value is not None + memory_value.text_cache = text_cache_layers[i] hidden_states, lbl_metadata_dict, kv_to_store = decoder_layer( hidden_states, @@ -1411,6 +1442,8 @@ def forward( before every decoder operation and restores the full layout on return. """ target_only_no_text = gen_only and bool(getattr(memory_value, "target_only_no_text", False)) + text_cache = getattr(memory_value, "text_cache", None) + cached_text = text_cache is not None and text_cache.hidden_output is not None layer_input = input layer_position_embeddings = packed_position_embeddings target_start = 0 @@ -1462,8 +1495,8 @@ def forward( ) # Pre-Attention layernorm - if target_only_no_text: - norm_und = get_und_seq(layer_input) # [0,hidden_size] + if target_only_no_text or cached_text: + norm_und = get_und_seq(layer_input) # [N_und,hidden_size], empty in target-only mode else: norm_und = _apply_request_rms_norm( self.input_layernorm, get_und_seq(layer_input), memory_value @@ -1563,8 +1596,10 @@ def forward( mlp_out_gen_seq = residual_gen + mlp_out_gen else: # STANDARD PATH: Process both und and gen tokens - ln_out_und = _apply_request_rms_norm( - self.post_attention_layernorm, residual_und, memory_value + ln_out_und = ( + residual_und + if cached_text + else _apply_request_rms_norm(self.post_attention_layernorm, residual_und, memory_value) ) # [N_und,hidden_size] ln_out_gen = _apply_request_rms_norm( self.post_attention_layernorm_moe_gen, residual_gen, memory_value @@ -1582,13 +1617,16 @@ def forward( if self._sample_lbl_und: und_sample_ids, und_num_samples = _get_local_sample_ids(layer_input, "und") - mlp_out_und, lbl_metadata_und = _run_mlp( - self.mlp, - ln_out_und, - und_token_mask, - sample_ids=und_sample_ids, - num_samples=und_num_samples, - ) + if cached_text: + mlp_out_und, lbl_metadata_und = ln_out_und, None # [N_und,hidden_size], unused caption residual + else: + mlp_out_und, lbl_metadata_und = _run_mlp( + self.mlp, + ln_out_und, + und_token_mask, + sample_ids=und_sample_ids, + num_samples=und_num_samples, + ) # mlp_out_und: [N_und,hidden_size], zero on the padding rows mlp_out_gen, lbl_metadata_gen = _run_mlp( self.mlp_moe_gen, @@ -1604,7 +1642,13 @@ def forward( if lbl_metadata_gen is not None: lbl_metadata_dict["gen"] = lbl_metadata_gen - mlp_out_und_seq = residual_und + mlp_out_und # [N_und,hidden_size] + if cached_text: + assert text_cache is not None and text_cache.hidden_output is not None + mlp_out_und_seq = text_cache.hidden_output # [N_und,hidden_size] + else: + mlp_out_und_seq = residual_und + mlp_out_und # [N_und,hidden_size] + if text_cache is not None: + text_cache.hidden_output = mlp_out_und_seq mlp_out_gen_seq = residual_gen + mlp_out_gen # [N_gen,hidden_size] layer_output = from_und_gen_splits(mlp_out_und_seq, mlp_out_gen_seq, layer_input) diff --git a/cosmos_framework/model/generator/mot/unified_mot_test.py b/cosmos_framework/model/generator/mot/unified_mot_test.py index 954c08d77..83b4ffb81 100644 --- a/cosmos_framework/model/generator/mot/unified_mot_test.py +++ b/cosmos_framework/model/generator/mot/unified_mot_test.py @@ -40,7 +40,7 @@ import math import os import subprocess -from collections.abc import Iterator +from collections.abc import Callable, Iterator from types import SimpleNamespace import pytest @@ -81,6 +81,7 @@ Qwen3VLMoeTextMLP, Qwen3VLMoeTextRMSNorm, ) +from cosmos_framework.model.generator.utils.memory import ARTextLayerCache, MemoryValue from cosmos_framework.data.generator.sequence_packing.runtime import ( SequencePack, from_und_gen_splits, @@ -94,6 +95,116 @@ # ----------------------------------------------------------------------------- +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("normalized_text_keys", [False, True]) +def test_ar_text_cache_reuses_layer_work_and_preserves_media_predictions(normalized_text_keys: bool) -> None: + """Media changes each step; caption projections/MLP run once per chunk/branch.""" + config = Qwen3VLTextConfig( + hidden_size=8, + intermediate_size=16, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=4, + num_hidden_layers=1, + ) + layer = MoTDecoderLayer( + config, + layer_idx=0, + layer_types=LayerTypes("qwen3_vl_dense"), + qk_norm_for_text=False, + qk_norm_for_diffusion=normalized_text_keys, + ).eval() + + def dense_attention( + q: SequencePack, + k: SequencePack, + v: SequencePack, + mask: object, + *, + memory_value: MemoryValue, + packed_key_states_normalized: SequencePack | None, + **kwargs: object, + ) -> tuple[SequencePack, None]: + cache = memory_value.text_cache + text_q, text_k, text_v = get_und_seq(q), get_und_seq(k), get_und_seq(v) # each [S_text,H,D] + if cache is not None and cache.attention_output is not None: + text_out = cache.attention_output # [S_text,H*D] + else: + text_out = ( + torch.nn.functional.scaled_dot_product_attention( + text_q.transpose(0, 1), + text_k.transpose(0, 1), + text_v.transpose(0, 1), + is_causal=True, + ) + .transpose(0, 1) + .flatten(-2) + ) # [S_text,H*D] + if cache is not None: + cache.attention_output = text_out + keys = k if packed_key_states_normalized is None else packed_key_states_normalized + media_k = torch.cat((get_und_seq(keys), get_gen_seq(keys))) # [S_text+S_gen,H,D] + media_v = torch.cat((text_v, get_gen_seq(v))) # [S_text+S_gen,H,D] + media_out = ( + torch.nn.functional.scaled_dot_product_attention( + get_gen_seq(q).transpose(0, 1), + media_k.transpose(0, 1), + media_v.transpose(0, 1), + ) + .transpose(0, 1) + .flatten(-2) + ) # [S_gen,H*D] + return from_und_gen_splits(text_out, media_out, q), None + + layer.self_attn.dispatch_attention_fn = dense_attention + counts: dict[str, int] = {"q_proj": 0, "mlp": 0} + + def count(name: str) -> Callable[[torch.nn.Module, tuple[object, ...], object], None]: + def hook(module: torch.nn.Module, args: tuple[object, ...], result: object) -> None: + counts[name] += 1 + + return hook + + layer.self_attn.q_proj.register_forward_hook(count("q_proj")) + layer.mlp.register_forward_hook(count("mlp")) + pack: SequencePack = { + "causal_seq": torch.randn(3, 8), + "full_only_seq": torch.randn(4, 8), # [S_text,D], [S_gen,D] + "is_sharded": False, + "_num_causal_tokens": 3, + "_num_full_tokens": 4, + "sample_offsets": torch.tensor([0, 7], dtype=torch.int32), # [B+1] + "max_sample_len": 7, + "max_causal_len": 3, + "max_full_len": 4, + "_causal_indices": torch.arange(3), # [S_text] + "_full_indices": torch.arange(3, 7), # [S_gen] + "_causal_seq_offsets": torch.tensor([0, 3], dtype=torch.int32), # [B+1] + "_full_only_seq_offsets": torch.tensor([0, 4], dtype=torch.int32), # [B+1] + } + cos = from_und_gen_splits(torch.ones(3, 4), torch.ones(4, 4), pack) # [S_text+S_gen,D_head] + sin = from_und_gen_splits(torch.zeros(3, 4), torch.zeros(4, 4), pack) # [S_text+S_gen,D_head] + cache = ARTextLayerCache() + memory = MemoryValue(text_cache=cache) + with torch.no_grad(): + layer(pack, object(), (cos, sin), memory_value=memory) + pack["full_only_seq"] = torch.randn(4, 8) # [S_gen,D], next denoising timestep + cached, _, cached_kv = layer(pack, object(), (cos, sin), memory_value=memory) + assert counts == {"q_proj": 1, "mlp": 1} + reference, _, reference_kv = layer(pack, object(), (cos, sin), memory_value=MemoryValue()) + torch.testing.assert_close(get_gen_seq(cached), get_gen_seq(reference), rtol=0, atol=0) + torch.testing.assert_close(get_und_seq(cached), get_und_seq(reference), rtol=0, atol=0) + assert cached_kv is not None and reference_kv is not None + for actual, expected in zip(cached_kv, reference_kv, strict=True): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + # Another chunk/CFG branch has a new cache, even at identical token shapes. + pack["causal_seq"] = torch.randn(3, 8) # [S_text,D] + fresh, _, _ = layer(pack, object(), (cos, sin), memory_value=MemoryValue(text_cache=ARTextLayerCache())) + uncached, _, _ = layer(pack, object(), (cos, sin), memory_value=MemoryValue()) + torch.testing.assert_close(get_gen_seq(fresh), get_gen_seq(uncached), rtol=0, atol=0) + + @pytest.mark.L0 @pytest.mark.CPU def test_decoder_target_only_path_skips_und_and_control_then_restores_full_layout() -> None: diff --git a/cosmos_framework/model/generator/mot/vae_load_balance_test.py b/cosmos_framework/model/generator/mot/vae_load_balance_test.py index 50d04e084..391a4c7da 100644 --- a/cosmos_framework/model/generator/mot/vae_load_balance_test.py +++ b/cosmos_framework/model/generator/mot/vae_load_balance_test.py @@ -415,6 +415,9 @@ def test_balanced_vae_encode_on_a_real_4gpu_node() -> None: world_size = int(os.environ["WORLD_SIZE"]) if not dist.is_initialized(): + # Same NVLS policy as the benchmark CLI and production launcher (NVLink SHARP fails on + # GB200); importing build_model no longer sets it process-wide. + os.environ.setdefault("NCCL_NVLS_ENABLE", "0") dist.init_process_group(backend="nccl") rank = dist.get_rank() torch.cuda.set_device(rank) diff --git a/cosmos_framework/model/generator/omni_mot_causal_model.py b/cosmos_framework/model/generator/omni_mot_causal_model.py index 90763ed99..d13fa45c5 100644 --- a/cosmos_framework/model/generator/omni_mot_causal_model.py +++ b/cosmos_framework/model/generator/omni_mot_causal_model.py @@ -28,37 +28,21 @@ from typing_extensions import override import cosmos_framework.model.generator.omni_mot_model as omni_mot_model_module -from cosmos_framework.configs.base.defaults.causal_flex_attention import CausalFlexAttentionConfig from cosmos_framework.configs.base.defaults.model_config import OmniMoTModelConfig +from cosmos_framework.configs.base.defaults.causal_flex_attention import CausalFlexAttentionConfig from cosmos_framework.configs.base.defaults.replay_attention import ( TEACHER_FORCING_KV_IMPLEMENTATIONS, TeacherForcingKVImplementation, TeacherForcingReplayPolicyConfig, ) from cosmos_framework.data.generator.augmentors.text_tokenizer import TEXT_SYSTEM_PROMPT_KEY -from cosmos_framework.data.generator.sequence_packing import ( - PackedSequence, - SequencePlan, - build_sequence_plans_from_data_batch, -) -from cosmos_framework.data.generator.sequence_packing.autoregressive import ( - pack_input_sequence_autoregressive, - pack_input_sequence_autoregressive_batch, - pack_multiview_action_temporal_causal, - resolve_text_system_prompt, -) -from cosmos_framework.data.generator.sequence_packing.modality import compute_text_split_length -from cosmos_framework.data.generator.sequence_packing.packers import ( - make_teacher_forcing_clean_pack, -) -from cosmos_framework.data.generator.sequence_packing.runtime import to_device_nonblocking from cosmos_framework.model.generator.attention_io_layout import AttentionIOLayout from cosmos_framework.model.generator.mot.causal_attention import dispatch_attention_with_memory +from cosmos_framework.model.generator.mot.causal_flex_attention import build_teacher_forcing_clean_target_token_indexes from cosmos_framework.model.generator.mot.causal_cosmos3_vfm_network import ( InteractiveCosmos3VFMNetwork, build_interactive_multiview_mask_items, ) -from cosmos_framework.model.generator.mot.causal_flex_attention import build_teacher_forcing_clean_target_token_indexes from cosmos_framework.model.generator.mot.maskless_attention import require_gapped_replay from cosmos_framework.model.generator.mot.post_saturation.installer import install_ar_post_saturation_mode from cosmos_framework.model.generator.mot.post_saturation.runtime import ( @@ -84,6 +68,7 @@ ) from cosmos_framework.model.generator.utils.kv_storage_backend import validate_kv_cache_dtype from cosmos_framework.model.generator.utils.memory import MemoryState +from cosmos_framework.model.generator.utils.model_construction import is_architecture_only_model_construction from cosmos_framework.model.generator.utils.multiview_ar import ( JointARCache, MultiviewTransferARBackend, @@ -93,6 +78,22 @@ sample_joint_transfer_ar, ) from cosmos_framework.model.generator.utils.nvfp4 import resolve_legacy_nvfp4_mode +from cosmos_framework.data.generator.sequence_packing import ( + PackedSequence, + SequencePlan, + build_sequence_plans_from_data_batch, +) +from cosmos_framework.data.generator.sequence_packing.autoregressive import ( + pack_input_sequence_autoregressive, + pack_input_sequence_autoregressive_batch, + pack_multiview_action_temporal_causal, + resolve_text_system_prompt, +) +from cosmos_framework.data.generator.sequence_packing.modality import compute_text_split_length +from cosmos_framework.data.generator.sequence_packing.packers import ( + make_teacher_forcing_clean_pack, +) +from cosmos_framework.data.generator.sequence_packing.runtime import to_device_nonblocking from cosmos_framework.utils.generator.data_batch import condition_frame_indexes_vision_from_batch _ARBranch = Literal["conditional", "unconditional"] @@ -101,7 +102,7 @@ # Framed-caption capacity chosen for long multi-caption rollouts with headroom; configurable for # other datasets, and used only for static-compile padding, never media positions. -_DEFAULT_ROLLING_STATIC_CAPTION_MAX_TOKENS: int = 576 +_DEFAULT_ROLLING_STATIC_CAPTION_MAX_TOKENS: int = 1024 class GaussianBellTrainTimeWeight: @@ -493,6 +494,7 @@ class OmniMoTCausalModel(OmniMoTModel): """ config: OmniMoTCausalModelConfig + ar_use_text_cache: bool = False # Set while a bidirectional MoBA step or dense MoBA inference runs. Read by # pre_noise_memory_hook and denoise. @@ -622,8 +624,9 @@ def build_net( or replay_policy.decomposed_temporal_window_seconds is not None ): raise ValueError("multiview_threeway_kv requires causal robot packs and same-instant decomposed attention.") - if maskless_replay: + if maskless_replay and not is_architecture_only_model_construction(): # Validate the execution device before network construction enters the meta device context. + # Checkpoint conversion only inspects architecture and never executes replay kernels. require_gapped_replay(torch.device(self.tensor_kwargs["device"])) # Robot multiview keeps the production three-way layout and adds only # one same-instant RGB varlen component in Interactive attention. @@ -2155,6 +2158,7 @@ def iter_samples_from_batch_autoregressive( vision_items = gen_data_clean.x0_tokens_vision if vision_items is None: raise ValueError(f"{mode} AR inference requires vision latents") + physical_view_ids = getattr(gen_data_clean, "vision_view_ids", None) is_transfer = mode == "video_transfer" # The RGB-control frontend sets this per-request marker after validating # edge/blur/depth/seg. Other modes and multiview dispatch never opt in. @@ -2500,6 +2504,7 @@ def iter_samples_from_batch_autoregressive( include_text = prefix_frame_idx == 0 self._seed_frame_into_kv_cache( frame_latent=prefix_vision_latent, + vision_view_ids=physical_view_ids[0] if physical_view_ids is not None else None, frame_idx=prefix_frame_idx, use_ar_rolling_path=use_ar_rolling_path, preserve_transfer_numerics=preserve_transfer_numerics, @@ -2579,6 +2584,7 @@ def report_history(next_frame_idx: int) -> None: frame_latent=control_latent, frame_idx=transfer_history_cache_idx, position_frame_idx=chunk_start, + vision_view_ids=physical_view_ids[0] if physical_view_ids is not None else None, dual_kv_cache=dual_kv_cache, dual_kv_cache_uncond=dual_kv_cache_uncond, cond_text_tokens=cond_text_tokens[0] if (cond_text_tokens and include_text) else None, @@ -2689,6 +2695,9 @@ def report_history(next_frame_idx: int) -> None: _cond_pack_text = cond_text_tokens[0] if (cond_text_tokens and include_text_in_pack) else None packed_seq = pack_input_sequence_autoregressive( vision_latent=curr_vision_latent, + vision_view_ids=( + physical_view_ids[generation_vision_item_idx] if physical_view_ids is not None else None + ), action_latent=curr_action_latent, text_tokens=_cond_pack_text, timestep=0.0, @@ -2717,6 +2726,9 @@ def report_history(next_frame_idx: int) -> None: _uncond_pack_text = uncond_text_tokens[0] if (uncond_text_tokens and include_text_in_pack) else None packed_seq_uncond = pack_input_sequence_autoregressive( vision_latent=curr_vision_latent, + vision_view_ids=( + physical_view_ids[generation_vision_item_idx] if physical_view_ids is not None else None + ), action_latent=curr_action_latent, text_tokens=_uncond_pack_text, timestep=0.0, @@ -2817,6 +2829,9 @@ def report_history(next_frame_idx: int) -> None: cache.gen_cache.select_history(_f, refresh_history) self._seed_frame_into_kv_cache( frame_latent=_frame_latent, + vision_view_ids=( + physical_view_ids[generation_vision_item_idx] if physical_view_ids is not None else None + ), frame_idx=transfer_history_cache_idx if is_transfer else _f, position_frame_idx=_f, dual_kv_cache=dual_kv_cache, @@ -2926,6 +2941,11 @@ def generate_joint_samples_from_batch_autoregressive( ) if self.parallel_dims is not None and self.parallel_dims.cp_enabled: seed = _broadcast_seed([seed], self.parallel_dims.cp_mesh.get_group(), self.parallel_dims.cp_rank)[0] + # Optional benchmark boundary after input encoding, before clean-prefix + # work and diffusion. The callback is scoped to this generation. + on_inputs_prepared: Callable[[], None] | None = data_batch.pop("_on_ar_inputs_prepared", None) + if on_inputs_prepared is not None: + on_inputs_prepared() return sample_joint_transfer_ar( self, plans=plans, @@ -3068,6 +3088,7 @@ def _seed_frame_into_kv_cache( transfer_history_sink_tokens: int = 0, transfer_history_max_tokens: int | None = None, batched_ar: bool = False, + vision_view_ids: torch.Tensor | None = None, # [V] identity of the clean RGB item being cached ) -> None: """Run one forward pass that writes ``frame_latent``'s K/V into ``dual_kv_cache[layer].gen_cache[frame_idx]``, under the strategy-appropriate @@ -3161,6 +3182,7 @@ def _build_pack( raw_action_dim = OmniMoTCausalModel._first_raw_action_dim(gen_data_clean) return pack_input_sequence_autoregressive( vision_latent=frame_in, + vision_view_ids=vision_view_ids, action_latent=curr_action_latent, text_tokens=text_tokens, timestep=timestep_val, diff --git a/cosmos_framework/model/generator/omni_mot_causal_model_test.py b/cosmos_framework/model/generator/omni_mot_causal_model_test.py index 683c7bf34..4c3030b7b 100644 --- a/cosmos_framework/model/generator/omni_mot_causal_model_test.py +++ b/cosmos_framework/model/generator/omni_mot_causal_model_test.py @@ -18,6 +18,7 @@ uninitialized. """ +from contextlib import nullcontext from types import MethodType, SimpleNamespace from typing import Literal, cast from unittest.mock import MagicMock, patch @@ -41,8 +42,8 @@ def test_clean_tf_cache_preserves_target_indexes_and_caption_layout( view_scoped: bool, detach_clean_kv: bool, maskless_replay: bool ) -> None: """Exercise the real mask builder and index helper through the model entry point.""" - from cosmos_framework.data.generator.sequence_packing.sequence import ModalityData, PackedSequence from cosmos_framework.model.generator.omni_mot_causal_model import OmniMoTCausalModel + from cosmos_framework.data.generator.sequence_packing.sequence import ModalityData, PackedSequence original_masks = [ torch.ones(4), # [V*T] @@ -295,8 +296,8 @@ def test_target_only_teacher_forcing_requires_one_logical_sample( num_items: int, item_counts: list[int] | None, sample_lens: list[int], valid: bool ) -> None: """Implicit single-target grouping is valid; multiple logical samples are not.""" - from cosmos_framework.data.generator.sequence_packing.sequence import ModalityData, PackedSequence from cosmos_framework.model.generator.omni_mot_causal_model import OmniMoTCausalModel + from cosmos_framework.data.generator.sequence_packing.sequence import ModalityData, PackedSequence model = MagicMock() model.config.video_temporal_causal = True @@ -409,13 +410,13 @@ def test_detached_target_only_cache_is_allowed_only_for_tfdcm_teacher( @pytest.mark.CPU def test_teacher_forcing_replay_policy_resolves_lazy_config_and_runs_validation() -> None: """LazyConfig must become the typed policy before model construction.""" + from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.configs.base.defaults.replay_attention import TeacherForcingReplayPolicyConfig from cosmos_framework.model.generator.omni_mot_causal_model import ( OmniMoTCausalModel, OmniMoTCausalModelConfig, _resolve_teacher_forcing_replay_policy, ) - from cosmos_framework.utils.lazy_config import LazyCall as L lazy_call = L(OmniMoTCausalModel)(config=OmniMoTCausalModelConfig(), _recursive_=False) lazy_policy = lazy_call.config.teacher_forcing_replay_policy @@ -449,12 +450,12 @@ def test_public_teacher_forcing_kv_implementation_survives_config_resolution( implementation: TeacherForcingKVImplementation, ) -> None: """Every public backend must pass both attrs validation and LazyConfig resolution.""" + from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.model.generator.omni_mot_causal_model import ( OmniMoTCausalModel, OmniMoTCausalModelConfig, _resolve_teacher_forcing_kv_implementation, ) - from cosmos_framework.utils.lazy_config import LazyCall as L config = OmniMoTCausalModelConfig(teacher_forcing_kv_implementation=implementation) lazy_model = L(OmniMoTCausalModel)(config=config, _recursive_=False) @@ -533,7 +534,8 @@ def test_multiview_replay_selection_rejects_non_replay_strategy(implementation: @pytest.mark.L0 @pytest.mark.CPU @pytest.mark.parametrize("causal_training_strategy", ["teacher_forcing", "teacher_forcing_dcm"]) -def test_maskless_replay_requires_backend_before_network_build(causal_training_strategy: str) -> None: +@pytest.mark.parametrize("device", ["cpu", "cuda:3"]) +def test_maskless_replay_requires_backend_before_network_build(causal_training_strategy: str, device: str) -> None: """Reject a missing replay backend before allocating the network or changing its config.""" from cosmos_framework.configs.base.defaults.replay_attention import TeacherForcingReplayPolicyConfig from cosmos_framework.model.generator import omni_mot_causal_model as causal_model_module @@ -542,7 +544,7 @@ def test_maskless_replay_requires_backend_before_network_build(causal_training_s model = object.__new__(OmniMoTCausalModel) torch.nn.Module.__init__(model) - model.tensor_kwargs = {"device": "cuda:3"} + model.tensor_kwargs = {"device": device} model.config = SimpleNamespace( causal_training_strategy=causal_training_strategy, teacher_forcing_kv_implementation="multiview_maskless_kv", @@ -564,7 +566,7 @@ def test_maskless_replay_requires_backend_before_network_build(causal_training_s ): model.build_net(torch.bfloat16) - require_replay.assert_called_once_with(torch.device("cuda:3")) + require_replay.assert_called_once_with(torch.device(device)) base_build.assert_not_called() assert model.config.video_temporal_causal is True assert model.config.joint_attn_implementation == "three_way" @@ -577,8 +579,9 @@ def test_maskless_replay_requires_backend_before_network_build(causal_training_s @pytest.mark.parametrize("implementation", ["multiview_flex_kv", "multiview_maskless_kv"]) @pytest.mark.parametrize("video_temporal_causal", [False, True]) @pytest.mark.parametrize("causal_training_strategy", ["teacher_forcing", "teacher_forcing_dcm"]) +@pytest.mark.parametrize("architecture_only", [False, True]) def test_teacher_forcing_build_keeps_student_and_teacher_backends_separate( - implementation: str, video_temporal_causal: bool, causal_training_strategy: str + implementation: str, video_temporal_causal: bool, causal_training_strategy: str, architecture_only: bool ) -> None: """Construct the selected student network while preserving bidirectional teacher settings.""" from cosmos_framework.configs.base.defaults.replay_attention import TeacherForcingReplayPolicyConfig @@ -586,6 +589,7 @@ def test_teacher_forcing_build_keeps_student_and_teacher_backends_separate( from cosmos_framework.model.generator.mot.causal_cosmos3_vfm_network import InteractiveCosmos3VFMNetwork from cosmos_framework.model.generator.omni_mot_causal_model import OmniMoTCausalModel from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel + from cosmos_framework.model.generator.utils.model_construction import architecture_only_model_construction replay_policy = TeacherForcingReplayPolicyConfig( control_visibility="causal", @@ -594,7 +598,7 @@ def test_teacher_forcing_build_keeps_student_and_teacher_backends_separate( ) model = object.__new__(OmniMoTCausalModel) torch.nn.Module.__init__(model) - model.tensor_kwargs = {"device": "cuda:3"} + model.tensor_kwargs = {"device": "cpu" if architecture_only else "cuda:3"} model.config = SimpleNamespace( causal_training_strategy=causal_training_strategy, teacher_forcing_kv_implementation=implementation, @@ -637,10 +641,11 @@ def _fake_base_build( with ( patch.object(OmniMoTModel, "build_net", side_effect=_fake_base_build), patch.object(causal_model_module, "require_gapped_replay") as require_replay, + architecture_only_model_construction() if architecture_only else nullcontext(), ): result = model.build_net(torch.bfloat16) - if use_maskless: + if use_maskless and not architecture_only: require_replay.assert_called_once_with(torch.device("cuda:3")) else: require_replay.assert_not_called() @@ -698,8 +703,8 @@ def _make_data() -> SimpleNamespace: @pytest.mark.parametrize("dropout_rate", [0.0, 1.0]) def test_target_only_replay_accepts_pack_after_control_dropout(self, dropout_rate: float) -> None: """Keep the grouping returned by real control dropout when validating replay.""" - from cosmos_framework.data.generator.sequence_packing.sequence import ModalityData, PackedSequence from cosmos_framework.model.generator.omni_mot_causal_model import OmniMoTCausalModel + from cosmos_framework.data.generator.sequence_packing.sequence import ModalityData, PackedSequence model = self._make_model(dropout_rate) model.config.video_temporal_causal = True @@ -935,8 +940,8 @@ def test_teacher_forcing_cp_local_kv_heads_match_ulysses_layout() -> None: def test_three_way_teacher_forcing_memory_state_does_not_require_flex_metadata() -> None: """Legacy replay defaults the Flex-only clean-memory capacity to zero.""" from cosmos_framework.configs.base.defaults.replay_attention import TeacherForcingReplayPolicyConfig - from cosmos_framework.data.generator.sequence_packing.sequence import ModalityData, PackedSequence from cosmos_framework.model.generator.omni_mot_causal_model import OmniMoTCausalModel + from cosmos_framework.data.generator.sequence_packing.sequence import ModalityData, PackedSequence model = object.__new__(OmniMoTCausalModel) torch.nn.Module.__init__(model) @@ -978,9 +983,9 @@ def test_multiview_clean_tf_cache_selects_target_tokens_and_preserves_condition_ ) -> None: """Both TF stages select real target tokens before the clean copy becomes fully conditioned.""" from cosmos_framework.configs.base.defaults.replay_attention import TeacherForcingReplayPolicyConfig - from cosmos_framework.data.generator.sequence_packing.sequence import ModalityData, PackedSequence from cosmos_framework.model.generator.omni_mot_causal_model import OmniMoTCausalModel from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean + from cosmos_framework.data.generator.sequence_packing.sequence import ModalityData, PackedSequence model = object.__new__(OmniMoTCausalModel) torch.nn.Module.__init__(model) @@ -1138,9 +1143,7 @@ def test_multiview_transfer_ar_mode_dispatches_to_specialized_iterator() -> None "sample_n_views": torch.tensor([2]), # [B] } - with patch( - "cosmos_framework.model.generator.omni_mot_causal_model.reset_ar_post_saturation_runtime_for_generation" - ): + with patch("cosmos_framework.model.generator.omni_mot_causal_model.reset_ar_post_saturation_runtime_for_generation"): outputs = list( OmniMoTCausalModel.iter_samples_from_batch_autoregressive( model, @@ -1416,11 +1419,11 @@ def test_multiview_transfer_ar_pack_sets_metadata_and_aligned_view_positions() - @pytest.mark.parametrize("condition_count", [0, 1, 5]) def test_multiview_transfer_prefill_is_clean_without_extending_conditioned_prefix(condition_count: int) -> None: """Recomputed RGB history omits diffusion embeddings without changing the rollout prefix.""" - from cosmos_framework.data.generator.sequence_packing import ModalityData from cosmos_framework.model.generator.utils.multiview_ar import ( MultiviewTransferARBackend, multiview_conditioned_prefix_length, ) + from cosmos_framework.data.generator.sequence_packing import ModalityData target_mask = (torch.arange(5) < condition_count).repeat(2).reshape(10, 1, 1) # [V*T,1,1] control_mask = torch.ones_like(target_mask) # [V*T,1,1] @@ -2376,6 +2379,30 @@ def _run( result = {"vision": torch.cat([out["vision"] for out in outputs], dim=2)} # [B,C,T,H,W] return result, mock_pack + @pytest.mark.L0 + @pytest.mark.CPU + @pytest.mark.parametrize("mode", ["text2video", "image2video", "video_transfer"]) + @pytest.mark.parametrize("view_id", [-1, 8]) + @pytest.mark.parametrize("guidance", [1.0, 3.0]) + def test_ar_preserves_view_ids_for_denoising_prefill_and_refresh( + self, mode: str, view_id: int, guidance: float + ) -> None: + model = self._make_model_mock() + data = self._make_gen_data(mode) + target_ids = torch.tensor([view_id]) # [V] + control_ids = torch.tensor([2]) # [V] + data.vision_view_ids = [control_ids, target_ids] if mode == "video_transfer" else [target_ids] + + _, build_pack = self._run(model, mode, gen_data=data, guidance=guidance) + + assert build_pack.call_count > 0 + for call in build_pack.call_args_list: + torch.testing.assert_close(call.kwargs["vision_view_ids"], target_ids) + assert model._seed_frame_into_kv_cache.call_count > 0 + for call in model._seed_frame_into_kv_cache.call_args_list: + is_control = mode == "video_transfer" and bool(call.kwargs["frame_latent"].eq(1).all()) + torch.testing.assert_close(call.kwargs["vision_view_ids"], control_ids if is_control else target_ids) + @pytest.mark.L0 @pytest.mark.CPU @pytest.mark.parametrize("mode", ["video_transfer", "text2video", "image2video", "forward_dynamics"]) @@ -4006,6 +4033,38 @@ def test_teacher_forcing_replay_policy_resolves_a_serialized_config_dict() -> No assert resolved_from_omegaconf == resolved +@pytest.mark.L0 +@pytest.mark.CPU +def test_teacher_forcing_replay_policy_resolves_lazycall_dictconfig() -> None: + """A LazyCall-style DictConfig (``_target_`` plus fields), as composed configs hold it. + + ``OmegaConf.to_object`` returned this node unchanged, so it reached the typed check as a + DictConfig and was rejected. Converting to plain containers takes the typed path. + """ + from cosmos_framework.configs.base.defaults.replay_attention import ( + TeacherForcingReplayPolicyConfig, + ) + from cosmos_framework.model.generator.omni_mot_causal_model import ( + _resolve_teacher_forcing_replay_policy, + ) + + policy = OmegaConf.create( + { + "_target_": "cosmos_framework.configs.base.defaults.replay_attention.TeacherForcingReplayPolicyConfig", + "control_visibility": "current", + "controls_read_strict_past_clean_rgb": True, + "clean_pass_causality": "chunk", + } + ) + + resolved = _resolve_teacher_forcing_replay_policy(policy) + + assert isinstance(resolved, TeacherForcingReplayPolicyConfig) + assert resolved.control_visibility == "current" + assert resolved.controls_read_strict_past_clean_rgb is True + assert resolved.clean_pass_causality == "chunk" + + @pytest.mark.L0 @pytest.mark.CPU def test_teacher_forcing_replay_policy_still_rejects_unknown_real_fields() -> None: @@ -4039,20 +4098,25 @@ def _forward_graph_model_mock() -> MagicMock: @pytest.mark.L0 @pytest.mark.CPU @pytest.mark.parametrize("frame_idx", [0, 2]) +@pytest.mark.parametrize("view_id", [None, -1, 8]) @pytest.mark.parametrize("preserve_transfer_numerics", [False, True]) -def test_seed_routes_static_frames_to_the_forward_cuda_graph(frame_idx: int, preserve_transfer_numerics: bool) -> None: +def test_seed_routes_static_frames_to_the_forward_cuda_graph( + frame_idx: int, view_id: int | None, preserve_transfer_numerics: bool +) -> None: """cuda_graph_scope="forward": cache idx >= 1 seeds replay a whole-forward graph, idx 0 stays eager.""" from cosmos_framework.model.generator.omni_mot_causal_model import OmniMoTCausalModel model = _forward_graph_model_mock() cond_pack = MagicMock(name="cond_pack") + ids = torch.tensor([view_id]) if view_id is not None else None # [V] or None with ( - patch(_PATCH_PACK, MagicMock(return_value=cond_pack)), + patch(_PATCH_PACK, MagicMock(return_value=cond_pack)) as build_pack, patch("cosmos_framework.model.generator.omni_mot_causal_model.run_ar_forward_cuda_graph") as run_graph, ): OmniMoTCausalModel._seed_frame_into_kv_cache( model, frame_latent=torch.zeros(1, 4, 1, 2, 2), # [B,C,T,H,W] + vision_view_ids=ids, frame_idx=frame_idx, dual_kv_cache=[MagicMock()], dual_kv_cache_uncond=None, @@ -4080,6 +4144,7 @@ def test_seed_routes_static_frames_to_the_forward_cuda_graph(frame_idx: int, pre transfer_history_sink_tokens=8, transfer_history_max_tokens=12, ) + assert build_pack.call_args.kwargs["vision_view_ids"] is ids if frame_idx == 0: run_graph.assert_not_called() model.denoise.assert_called_once() diff --git a/cosmos_framework/model/generator/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py index f438aedc0..e46666a0f 100644 --- a/cosmos_framework/model/generator/omni_mot_model.py +++ b/cosmos_framework/model/generator/omni_mot_model.py @@ -82,6 +82,7 @@ prepare_camera_relative_poses, prepare_prope_cross_view_poses, ) +from cosmos_framework.model.generator.utils.component_unipc import ComponentSchedule, sample_component_unipc from cosmos_framework.model.generator.utils.data_and_condition import ( GenerationDataClean, GenerationDataNoised, @@ -1014,8 +1015,14 @@ def _pack_input_sequence( """ assert self.tokenizer_vision_gen is not None multiview_action_conditioning = getattr(self.config, "multiview_action_conditioning", None) - # Check source clocks before packing's fallback hides missing FPS. - if self.config.joint_attn_implementation == "multiview": + # Check source clocks before packing's fallback hides missing FPS. The + # attention path only applies multiview windows to packs with explicit layouts. + if self.config.joint_attn_implementation == "multiview" and ( + gen_data_clean.num_views_per_vision_item is not None + or gen_data_clean.num_lidar_items_per_sample is not None + or gen_data_clean.num_radar_items_per_sample is not None + or (multiview_action_conditioning is not None and gen_data_clean.num_vision_items_per_sample is not None) + ): same_view_window = self.config.multiview_attention.mask.same_view_temporal_window validate_cross_view_fps( sequence_plans, @@ -2353,8 +2360,9 @@ def _get_train_noise_level_vision( num_tokens: Number of tokens for each sample (before 2x2 merge). Needed for dynamic shift. sample_is_image: Per-sample image identities for mixed image/video schedules. sequence_plans: Effective per-sample modality metadata after CP input selection. - Batches with LiDAR and no RGB use shift_lidar when configured. RGB and joint - RGB/LiDAR batches share the vision schedule. None retains the vision schedule. + Samples with LiDAR and no RGB use shift_lidar when configured, including in + mixed batches. RGB and joint RGB/LiDAR samples retain the vision schedule. + None retains the vision schedule. Returns: (timesteps, sigmas): Both [B,1] for TF/base, or [B,T_max] for diffusion_forcing. @@ -2401,17 +2409,21 @@ def _get_train_noise_level_vision( rf_config = self.config.rectified_flow_training_config shift_image = getattr(rf_config, "shift_image", None) shift_config = shift_image if is_image_batch and shift_image is not None else rf_config.shift - if ( - sequence_plans - and getattr(rf_config, "shift_lidar", None) is not None - and all(plan.has_lidar and not plan.has_vision for plan in sequence_plans) - ): - shift_config = rf_config.shift_lidar - if not isinstance(shift_config, int) or shift_config <= 0: + shift_lidar = getattr(rf_config, "shift_lidar", None) + lidar_only_indexes = ( + [i for i, plan in enumerate(sequence_plans) if plan.has_lidar and not plan.has_vision] + if sequence_plans is not None and shift_lidar is not None + else [] + ) + if lidar_only_indexes: + if not isinstance(shift_lidar, int) or shift_lidar <= 0: raise ValueError("shift_lidar must be a positive integer") + if len(lidar_only_indexes) == batch_size: + # Pure LiDAR batches do not need vision resolution or token-count metadata. + shift_config = shift_lidar if isinstance(shift_config, int): # Int-based shift: use directly for all samples - shifts = torch.full((batch_size,), shift_config, dtype=torch.float32) + shifts = torch.full((batch_size,), shift_config, dtype=torch.float32) # [B] else: # Convert to plain dict to avoid traceback-based memory leaks when GC is disabled # (OmegaConf's `in` operator uses exception control flow internally). @@ -2455,7 +2467,11 @@ def _get_train_noise_level_vision( f"Resolution '{resolution}' not found in shift dict. Available resolutions: {list(shift_dict.keys())}" ) shifts_list.append(shift_dict[resolution]) - shifts = torch.tensor(shifts_list, dtype=torch.float32) + shifts = torch.tensor(shifts_list, dtype=torch.float32) # [B] + + if shift_lidar is not None and lidar_only_indexes: + # Override before frame expansion while preserving the sampling call and RNG order. + shifts[lidar_only_indexes] = shift_lidar # [N_lidar_only] # Sample noise times: B×T_max for DF (one per video latent frame), B×1 for base/TF if self.config.causal_training_strategy == "diffusion_forcing": @@ -3489,6 +3505,7 @@ def _get_velocity( packed_sequence_template: PackedSequence | None = None, memory: MemoryState | None = None, has_noisy_actions: bool, + action_timestep: torch.Tensor | None = None, ) -> list[torch.Tensor]: """ Compute velocity prediction for a single sampling step. @@ -3693,6 +3710,11 @@ def _get_velocity( timestep, ) + if action_timestep is not None: + if not has_noisy_actions or packed_sequence.action is None: + raise ValueError("Independent action timestep requires noisy action tokens") + self._copy_timestep_to_template(packed_sequence.action.timesteps, action_timestep) + # --- Network forward --- out = self.denoise( net=net, @@ -3933,6 +3955,114 @@ def _run_classifier_free_guidance( else: return other_v_list, v_list + def _resolve_action_vision_shift(self, data_batch: dict, gen_data_clean: GenerationDataClean) -> float: + """Resolve one shared vision shift from the training configuration.""" + rf_config = self.config.rectified_flow_training_config + shift_image = getattr(rf_config, "shift_image", None) + vision_shift_config = ( + shift_image if gen_data_clean.is_image_batch and shift_image is not None else rf_config.shift + ) + if isinstance(vision_shift_config, (int, float)): + return float(vision_shift_config) + + image_sizes = data_batch["image_size"] + sizes = ( + image_sizes.reshape(-1, 4) + if isinstance(image_sizes, torch.Tensor) + else torch.cat([size.reshape(-1, 4) for size in image_sizes]) + ) + resolutions = {get_vision_data_resolution((int(size[0]), int(size[1]))) for size in sizes} + shifts = {float(vision_shift_config[resolution]) for resolution in resolutions} + if len(shifts) != 1: + raise ValueError("Independent inference batches must share a vision shift") + return shifts.pop() + + def _action_inference_vision_lengths( + self, + gen_data_clean: GenerationDataClean, + sequence_plans: list[SequencePlan], + initial_noise: list[torch.Tensor], + has_noisy_actions: bool, + ) -> list[int]: + """Validate RGB/action packing and count vision elements for each sample.""" + if ( + not has_noisy_actions + or gen_data_clean.x0_tokens_lidar is not None + or gen_data_clean.x0_tokens_radar is not None + or gen_data_clean.x0_tokens_sound is not None + or not all(plan.has_action for plan in sequence_plans) + ): + raise ValueError("Independent action inference requires an RGB/action layout for every sample") + if not gen_data_clean.x0_tokens_vision or gen_data_clean.x0_tokens_action is None: + raise ValueError("Independent action inference requires RGB and action components") + vision_counts = gen_data_clean.num_vision_items_per_sample or [1] * len(initial_noise) + if ( + len(vision_counts) != len(initial_noise) + or len(gen_data_clean.x0_tokens_action) != len(initial_noise) + or sum(vision_counts) != len(gen_data_clean.x0_tokens_vision) + ): + raise ValueError("Independent inference vision layout differs from sample metadata") + vision_lengths: list[int] = [] + vision_offset = 0 + for sample_index, count in enumerate(vision_counts): + vision_length = sum( + item.numel() for item in gen_data_clean.x0_tokens_vision[vision_offset : vision_offset + count] + ) + vision_offset += count + if ( + initial_noise[sample_index].numel() + != vision_length + gen_data_clean.x0_tokens_action[sample_index].numel() + ): + raise ValueError("Independent inference requires flat [vision | action] samples") + vision_lengths.append(vision_length) + return vision_lengths + + def _action_inference_schedule( + self, + *, + data_batch: dict, + gen_data_clean: GenerationDataClean, + sequence_plans: list[SequencePlan], + initial_noise: list[torch.Tensor], + has_noisy_actions: bool, + shift: float, + explicit_vision_shift: bool, + action_shift: float | None, + scheduler_type: str, + sampler: Any, + velocity_postprocess: Any, + ) -> tuple[float, ComponentSchedule | None]: + """Resolve the opt-in RGB/action schedule without changing legacy callers.""" + schedule: ComponentSchedule | None = None + rf_config = self.config.rectified_flow_training_config + if action_shift is None and has_noisy_actions and getattr(rf_config, "independent_action_schedule", False): + action_shift = rf_config.shift_action + if action_shift is None: + raise ValueError("Independent action inference requires an explicit shift_action") + if ( + not explicit_vision_shift + and gen_data_clean.x0_tokens_vision + and (action_shift is not None or getattr(rf_config, "independent_action_schedule", False)) + ): + shift = self._resolve_action_vision_shift(data_batch, gen_data_clean) + if action_shift is not None: + if scheduler_type != "unipc" or not isinstance(sampler, UniPCSampler): + raise ValueError("Independent action inference requires UniPC") + if velocity_postprocess is not None: + raise ValueError("Independent action inference does not support velocity postprocessing") + vision_lengths = self._action_inference_vision_lengths( + gen_data_clean, sequence_plans, initial_noise, has_noisy_actions + ) + schedule = ComponentSchedule( + shifts=(shift, float(action_shift)), + lengths=[ + (length, noise.numel() - length) + for length, noise in zip(vision_lengths, initial_noise, strict=True) + ], + ) + log.info(f"Using modality schedules: vision_shift={shift}, action_shift={action_shift}") + return shift, schedule + @torch.no_grad() def generate_samples_from_batch( self, @@ -3952,6 +4082,7 @@ def generate_samples_from_batch( has_negative_prompt: bool = False, num_steps: int = 35, shift: float | None = None, + action_shift: float | None = None, sigma_max: float = 80.0, skip_text_tokens_for_cfg: bool = False, normalize_cfg: bool = False, @@ -4053,7 +4184,7 @@ def generate_samples_from_batch( - "vision": List of vision latent tensors (one per sample, variable shapes) - "action": List of external-space action tensors (only present when the packed sequence marks action tokens as - denoising targets, i.e. WAM / inverse dynamics. Plans whose + denoising targets, i.e. Policy / WAM / inverse dynamics. Plans whose actions are pure conditioning — forward dynamics, temporal-causal action conditioning — generate no actions and return no key) @@ -4062,6 +4193,7 @@ def generate_samples_from_batch( ValueError: If the seed is a single integer. This is not supported anymore: `seed` must be a list of integers, one for each sample. """ + explicit_vision_shift = shift is not None if shift is None: shift = 5.0 rf_config = getattr(self.config, "rectified_flow_training_config", None) @@ -4317,7 +4449,9 @@ def generate_samples_from_batch( elif guidance != 1.0 or velocity_postprocess_builder is not None: uncond_text_kv_cache = self._make_inference_text_kv_cache(target_net) - def velocity_fn(noise_x: list[torch.Tensor], timestep: torch.Tensor) -> list[torch.Tensor]: + def velocity_fn( + noise_x: list[torch.Tensor], timestep: torch.Tensor, action_timestep: torch.Tensor | None = None + ) -> list[torch.Tensor]: # len(noise_x) == B, noise_x[i] is shape (D) # timestep is shape (B, 1) torch.compiler.cudagraph_mark_step_begin() @@ -4327,6 +4461,8 @@ def velocity_fn(noise_x: list[torch.Tensor], timestep: torch.Tensor) -> list[tor # Expand timestep to (B, 1) timestep = timestep.repeat(len(noise_x), 1) # [B,1] + if action_timestep is not None: + action_timestep = action_timestep.repeat(len(noise_x), 1) # [B,1] def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool) -> list[torch.Tensor]: nonlocal uncond_text_kv_cache @@ -4348,6 +4484,7 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool) -> list net=net, noise_x=noise_x, timestep=timestep, + action_timestep=action_timestep, text_tokens=tokens, sequence_plans=sequence_plans, gen_data_clean=gen_data_clean, @@ -4417,6 +4554,11 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool) -> list net=net, noise_x=noise_x + noise_x, timestep=torch.cat([timestep, timestep], dim=0), + action_timestep=( + torch.cat([action_timestep, action_timestep], dim=0) + if action_timestep is not None + else None + ), text_tokens=list(cond_tokens) + list(uncond_tokens), sequence_plans=sequence_plans_doubled, gen_data_clean=gen_data_clean_doubled, @@ -4503,13 +4645,6 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool) -> list _max_num_steps = num_steps _extra_num_steps = _max_num_steps - num_steps - if isinstance(sampler, FixedStepSampler): - log.info(f"Using sampler: FixedStep (t_list={sampler.t_list}, sample_type={sampler.sample_type})") - elif scheduler_type == "unipc": - log.info(f"Using sampler: UniPC (shift={shift}, num_steps={num_steps})") - else: - log.info(f"Using sampler: EDM (sigma_max={sigma_max}, num_steps={num_steps})") - fixed_step_sampler_kwargs = {} if isinstance(sampler, FixedStepSampler): fixed_step_sampler_kwargs = { @@ -4517,6 +4652,27 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool) -> list "condition_mask": condition_mask, } + shift, component_schedule = self._action_inference_schedule( + data_batch=data_batch, + gen_data_clean=gen_data_clean, + sequence_plans=sequence_plans, + initial_noise=initial_noise, + has_noisy_actions=has_noisy_actions, + shift=shift, + explicit_vision_shift=explicit_vision_shift, + action_shift=action_shift, + scheduler_type=scheduler_type, + sampler=sampler, + velocity_postprocess=velocity_postprocess, + ) + + if isinstance(sampler, FixedStepSampler): + log.info(f"Using sampler: FixedStep (t_list={sampler.t_list}, sample_type={sampler.sample_type})") + elif scheduler_type == "unipc": + log.info(f"Using sampler: UniPC (shift={shift}, num_steps={num_steps})") + else: + log.info(f"Using sampler: EDM (sigma_max={sigma_max}, num_steps={num_steps})") + # Explicit sampler steps drive both cache indexing and mixed precision. # Select W8A16/W8A8 once per sampler step; reset (trace + staging cleanup) # when the request ends, including on error. @@ -4528,17 +4684,35 @@ def _step_callback(step_index: int, step_count: int) -> None: if _mixed_precision_runtime is not None: _mixed_precision_runtime.set_step(step_index, step_count) + def _sample_latents( + noise: list[torch.Tensor], # B x [N_i] + steps: int, + step_callback: Callable[[int, int], None] | None = None, + ) -> list[torch.Tensor]: # B x [N_i] + if component_schedule is not None: + # Modality order and timestep binding belong to inference; + # the coordinator and native UniPC schedulers are generic. + return sample_component_unipc( + lambda values, timesteps: velocity_fn(values, timesteps[0], timesteps[1]), + noise, + cfg=sampler.cfg, + schedule=component_schedule, + num_steps=steps, + step_callback=step_callback, + ) + return sampler( + velocity_fn, + noise, + num_steps=steps, + shift=shift, + seed=seed, + step_callback=step_callback, + **fixed_step_sampler_kwargs, + ) + try: if isinstance(sampler, FixedStepSampler) or scheduler_type == "unipc": - latents = sampler( - velocity_fn, - initial_noise, - num_steps=num_steps, - shift=shift, - seed=seed, - step_callback=_step_callback, - **fixed_step_sampler_kwargs, - ) + latents = _sample_latents(initial_noise, num_steps, _step_callback) # B x [N_i] if _extra_num_steps > 0: # Unregistered padding calls bypass caching on every shard rank. cache_step_index = None @@ -4555,14 +4729,7 @@ def _step_callback(step_index: int, step_count: int) -> None: # Deterministic cheap padding: the FSDP-alignment dummy # call must not inherit the last real step's W8A16. _mixed_precision_runtime.set_base_precision() - _ = sampler( - velocity_fn, - latents, - num_steps=_extra_num_steps, - shift=shift, - seed=seed, - **fixed_step_sampler_kwargs, - ) + _ = _sample_latents(latents, _extra_num_steps) # B x [N_i] else: # EDM Sampler chunk_sizes = [_x.shape[0] for _x in initial_noise] diff --git a/cosmos_framework/model/generator/parallelize_vlm.py b/cosmos_framework/model/generator/parallelize_vlm.py index 6951e6dad..45c18ff04 100644 --- a/cosmos_framework/model/generator/parallelize_vlm.py +++ b/cosmos_framework/model/generator/parallelize_vlm.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -"""Activation checkpointing, ``torch.compile`` and FSDP2 for Cosmos3 VLM ``HFModel``s. +"""Context-parallel setup, activation checkpointing, compile and FSDP2 for Cosmos3 VLMs. Hosts the single VLM-specific ``parallelize`` entry point used by ``vlm_model.VLMModel._init_vlm``. Lives under ``cosmos_framework/model/generator/`` @@ -18,6 +18,7 @@ from typing import Any import torch +import torch.distributed as dist import torch.nn as nn from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( checkpoint_wrapper as ptd_checkpoint_wrapper, @@ -28,15 +29,16 @@ create_selective_checkpoint_contexts, ) +from cosmos_framework.utils import log 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.parallelism import ( PRECISION_TO_TORCH_DTYPE, ParallelismConfig, ) +from cosmos_framework.model.generator.context_parallel import ContextParallelBatchPolicy, shard_decoder_inputs from cosmos_framework.model.generator.hf_model import HFModel from cosmos_framework.model.generator.reasoner.qwen35_caption import qwen35_fp32_decay_modules -from cosmos_framework.utils import log from cosmos_framework.utils.generator.parallelism import ParallelDims, fsdp_mesh # (parent, attribute name, module) — the module currently registered at that slot, plus what @@ -45,6 +47,64 @@ _BlockSlot = tuple[nn.Module, str, nn.Module] +def apply_cp(model: HFModel, group: dist.ProcessGroup, *, allow_multimodal: bool = False) -> None: + """Attach shared CP to explicitly bound native decoder attention and inputs. + + Configure before activation checkpointing and FSDP wrap native modules. + Native fusion, RoPE and attention kernels remain owned by model integrations. + WORLD FSDP averaging and the CP-scaled objective aggregate replicated vision + gradients; adding a separate vision all-reduce would double count. + """ + spec = model.context_parallel_spec() + if not spec.fused_ce_enabled: + raise ValueError("VLM CP requires fused weighted CE") + if getattr(model.model, "_cosmos_cp_group", None) is not None: + raise ValueError("VLM CP is already configured") + config = spec.text_config + if spec.require_default_rope and (config.rope_scaling or {}).get("rope_type", "default") != "default": + raise NotImplementedError("This native CP integration requires default RoPE") + size = dist.get_world_size(group) + if config._attn_implementation != "cosmos": + raise ValueError("VLM CP requires cosmos attention") + if size not in (2, 4): + raise NotImplementedError("VLM CP currently supports only CP=2 or CP=4") + if config.num_attention_heads % size or config.num_key_value_heads % size: + raise ValueError("CP must divide both query and KV head counts") + if getattr(config, "attention_dropout", 0.0) != 0.0: + raise ValueError("VLM CP requires attention_dropout=0") + if not spec.attention_modules: + raise ValueError("VLM CP requires explicit native decoder attention modules") + + if spec.keep_full_inputs: + + def shard_native_decoder_inputs( + module: nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any] + ) -> tuple[tuple[Any, ...], dict[str, Any]]: + if args or kwargs.get("input_ids") is not None: + raise ValueError("CP expects native fused embeddings at decoder entry") + if spec.decoder_packing_key is not None and kwargs.get(spec.decoder_packing_key) is None: + raise ValueError("CP requires global native packed boundaries") + if kwargs.get("past_key_values") is not None or kwargs.get("use_cache"): + raise ValueError("CP does not support a KV cache") + output = shard_decoder_inputs(kwargs, group, position_axes=spec.position_axes) + if output.get("cache_position") is not None: + raise ValueError("CP does not support explicit cache positions") + return args, output + + spec.decoder.register_forward_pre_hook(shard_native_decoder_inputs, with_kwargs=True) + setattr(model.model, "_cosmos_cp_group", group) + if allow_multimodal and spec.fusion_module is not None: + setattr(spec.fusion_module, "_cosmos_multimodal_cp_group", group) + for attention in spec.attention_modules: + setattr(attention, "_cosmos_cp_group", group) + model.context_parallel_batch_policy = ContextParallelBatchPolicy( + position_axes=spec.position_axes, + allow_multimodal=allow_multimodal, + keep_full_inputs=spec.keep_full_inputs, + requires_packed=spec.requires_packed, + ) + + def _collect_repeated_blocks(inner: nn.Module) -> tuple[list[_BlockSlot], set[str]]: """Collect the repeated transformer blocks by their ORIGINAL type name. @@ -129,14 +189,14 @@ def _resolve_reduce_dtype(parallelism_config: ParallelismConfig) -> torch.dtype: def _collect_ignored_audio_params(model: HFModel) -> set[nn.Parameter]: """Return immutable audio parameters that must stay outside FSDP. - The selected audio encoder is always a frozen, artifact-backed feature extractor. - The projector remains part of the raw HF root (and therefore FSDP-managed) + Frozen artifact-backed audio parameters stay replicated, while trainable + encoder parameters remain FSDP-managed. The projector stays in the raw HF root even when it is frozen, matching the existing vision merger/projector lifecycle and keeping checkpoint tensor layouts uniform. """ if not model.sound_und: return set() - return set(model.model.sound_und_model.encoder.parameters()) + return {parameter for parameter in model.model.sound_und_model.encoder.parameters() if not parameter.requires_grad} def _compile_kwargs(compile_config: CompileConfig) -> dict[str, Any]: @@ -513,12 +573,15 @@ def parallelize( precision: str, activation_checkpointing: ActivationCheckpointingConfig, compile_config: CompileConfig | None = None, + allow_multimodal_cp: bool = False, ) -> None: """Optimize an HFModel in place: activation checkpointing, ``torch.compile``, FSDP2. Mirrors ``parallelize_unified_mot.parallelize``, including its pass ORDER, which is load-bearing rather than incidental: + 0. :func:`apply_cp` attaches shared CP to native decoder interfaces when enabled, + preserving the integration's existing input-sharding stage. 1. :func:`apply_ac` collects the repeated blocks while their original type names are intact and swaps each for a ``CheckpointWrapper``. 2. :func:`apply_compile` wraps those in ``torch.compile``, so the ``checkpoint`` call @@ -550,7 +613,11 @@ def parallelize( ``preserve_rng_state``, ``determinism_check``). compile_config: Optional ``CompileConfig``; ``None``/``enabled=False`` skips compile. + allow_multimodal_cp: Enable replicated native media fusion for integrations + with an opt-in multimodal path. """ + if parallel_dims.cp_enabled: + apply_cp(model, parallel_dims.cp_mesh.get_group(), allow_multimodal=allow_multimodal_cp) slots = apply_ac(model, activation_checkpointing) slots = apply_compile(slots, compile_config) apply_fsdp(model, slots, parallel_dims, parallelism_config, precision) diff --git a/cosmos_framework/model/generator/reasoner/audio/registry.py b/cosmos_framework/model/generator/reasoner/audio/registry.py index d26beecd9..3361de6a3 100644 --- a/cosmos_framework/model/generator/reasoner/audio/registry.py +++ b/cosmos_framework/model/generator/reasoner/audio/registry.py @@ -34,7 +34,12 @@ def _build_parakeet_processor() -> Any: return ParakeetAudioProcessor() -def _build_qwen3_omni_thinking_model(projection_hidden_size: int, out_hidden_size: int) -> Any: +def _build_qwen3_omni_thinking_variant( + projection_hidden_size: int, + out_hidden_size: int, + *, + encoder_output_mode: str, +) -> Any: from cosmos_framework.model.generator.reasoner.qwen3_omni_audio.configuration_qwen3_omni_audio import ( Qwen3OmniAudioConfig, ) @@ -44,10 +49,27 @@ def _build_qwen3_omni_thinking_model(projection_hidden_size: int, out_hidden_siz Qwen3OmniAudioConfig( projection_hidden_size=projection_hidden_size, out_hidden_size=out_hidden_size, + encoder_output_mode=encoder_output_mode, ) ) +def _build_qwen3_omni_thinking_model(projection_hidden_size: int, out_hidden_size: int) -> Any: + return _build_qwen3_omni_thinking_variant( + projection_hidden_size, + out_hidden_size, + encoder_output_mode="native", + ) + + +def _build_qwen3_omni_thinking_pre_projection_model(projection_hidden_size: int, out_hidden_size: int) -> Any: + return _build_qwen3_omni_thinking_variant( + projection_hidden_size, + out_hidden_size, + encoder_output_mode="pre_projection", + ) + + def _build_qwen3_omni_thinking_processor() -> Any: from cosmos_framework.data.generator.processors.qwen3_omni_audio_processor import Qwen3OmniAudioProcessor @@ -63,6 +85,10 @@ def _build_qwen3_omni_thinking_processor() -> Any: build_model=_build_qwen3_omni_thinking_model, build_processor=_build_qwen3_omni_thinking_processor, ), + "qwen3_omni_thinking_pre_projection": AudioEncoderBackend( + build_model=_build_qwen3_omni_thinking_pre_projection_model, + build_processor=_build_qwen3_omni_thinking_processor, + ), } diff --git a/cosmos_framework/model/generator/reasoner/audio/utils.py b/cosmos_framework/model/generator/reasoner/audio/utils.py index 69361520c..b1dc79051 100644 --- a/cosmos_framework/model/generator/reasoner/audio/utils.py +++ b/cosmos_framework/model/generator/reasoner/audio/utils.py @@ -134,7 +134,7 @@ def merge_projected_audio_embeddings( "audio_token_lengths", "audio_sample_token_lengths", ) -_SUPPORTED_AUDIO_REASONER_TYPES = frozenset({"nemotron_siglip2", "qwen3_vl"}) +_SUPPORTED_AUDIO_REASONER_TYPES = frozenset({"nemotron_siglip2", "qwen3_vl", "qwen3_vl_moe"}) _PATCH_MARKER = "_cosmos_reasoner_audio_forward_patch" _TIED_OUTPUT_PATCH_MARKER = "_cosmos_reasoner_boundary_only_output_patch" @@ -237,6 +237,22 @@ def _projector_dummy_zero(model: torch.nn.Module) -> torch.Tensor | None: return model.sound_und_model.projector(dummy_input).sum() * 0 +def _encoder_dummy_zero(model: Any) -> torch.Tensor | None: # returns scalar or None + """Exercise a trainable audio tower before the vision/LLM forward on non-audio batches.""" + # Raw HF models and registered audio backends expose these attributes dynamically. + encoder = model.sound_und_model.encoder + if not model.training or not torch.is_grad_enabled() or not any(p.requires_grad for p in encoder.parameters()): + return None + reference = next(encoder.parameters()) # [...] + lengths = torch.tensor([100], device=reference.device, dtype=torch.long) # [1] + features = torch.zeros( + 1, 100, encoder.config.num_mel_bins, device=reference.device, dtype=reference.dtype + ) # [1,100,M] + token_lengths = encoder.get_output_lengths(lengths) # [1] + projected, _ = model.sound_und_model(features, lengths, token_lengths) # [1,T,D], [1] + return projected.sum() * 0 # scalar + + def _connect_projector_zero_to_logits(output: Any, projector_zero: torch.Tensor, labels: Any) -> Any: """Connect a no-op projector graph to logits for mixed-modality FSDP ranks.""" @@ -272,7 +288,8 @@ def patch_reasoner_audio_forward( The patch deliberately wraps the raw conditional-generation model rather than its text tower. This keeps the audio projector owned by (and executed inside) the same FSDP root as the native vision merger. With audio absent, - the original forward receives the exact original arguments. With audio + the original forward receives the original arguments, with a zero-valued + audio dependency during training when the audio encoder is trainable. With audio present, native multimodal positions are computed from the untouched token IDs before only audio placeholder embeddings are replaced; the original forward then performs its normal image/video encoding and masked scatter. @@ -336,6 +353,23 @@ def patched_forward(self: torch.nn.Module, *args: Any, **kwargs: Any) -> Any: raise ValueError("Reasoner input contains audio placeholders, but audio inputs were not provided") if normalized_trainable_token_ids and self.training and torch.is_grad_enabled(): raise ValueError("Boundary-token-only training requires audio inputs on every training batch") + encoder_zero = _encoder_dummy_zero(self) # scalar or None + if encoder_zero is not None: + # Keep the audio graph ahead of vision/LLM backward, without changing + # native token IDs, multimodal positions, or another model instance. + def connect(_module: torch.nn.Module, _inputs: Any, embeddings: torch.Tensor) -> torch.Tensor: + # embeddings: [B,N,D], returns [B,N,D] + return embeddings + encoder_zero.to(device=embeddings.device, dtype=embeddings.dtype) # [B,N,D] + + if inputs_embeds is not None: + connected = connect(self, (), inputs_embeds) # [B,N,D] + _set_bound_argument(bound_arguments, original_signature, "inputs_embeds", connected) + return original_forward(*bound_arguments.args, **bound_arguments.kwargs) + hook = getattr(self, "get_input_embeddings")().register_forward_hook(connect) + try: + return original_forward(*args, **kwargs) + finally: + hook.remove() output = original_forward(*args, **kwargs) projector_zero = _projector_dummy_zero(self) if projector_zero is None: diff --git a/cosmos_framework/model/generator/reasoner/audio/utils_test.py b/cosmos_framework/model/generator/reasoner/audio/utils_test.py index 121a91a68..45d7d9717 100644 --- a/cosmos_framework/model/generator/reasoner/audio/utils_test.py +++ b/cosmos_framework/model/generator/reasoner/audio/utils_test.py @@ -5,6 +5,7 @@ import inspect from dataclasses import dataclass +from types import SimpleNamespace from typing import Any import pytest @@ -40,6 +41,10 @@ def __init__(self) -> None: super().__init__() self.scale = nn.Parameter(torch.ones(_HIDDEN_SIZE), requires_grad=False) self.forward_calls = 0 + self.config: SimpleNamespace = SimpleNamespace(num_mel_bins=_HIDDEN_SIZE) + + def get_output_lengths(self, lengths: torch.Tensor) -> torch.Tensor: # [B] -> [B] + return lengths def forward( self, @@ -482,3 +487,52 @@ def test_no_audio_forward_is_unchanged_and_connects_only_trainable_projector() - assert projector_grad is not None assert torch.count_nonzero(projector_grad) == 0 assert model.sound_und_model.encoder.scale.grad is None + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("use_inputs_embeds", [False, True]) +def test_trainable_audio_encoder_connects_on_vision_batches_without_global_patches(use_inputs_embeds: bool) -> None: + model = _FakeRawReasoner() + model.sound_und_model.encoder.requires_grad_(True) + other_model = _FakeRawReasoner() + for instance in (model, other_model): + patch_reasoner_audio_forward(instance, audio_token_id=_AUDIO_TOKEN_ID, model_type="qwen3_vl") + input_ids = torch.tensor([[1, _IMAGE_TOKEN_ID, 3]]) # [1,3] + kwargs: dict[str, Any] = {"input_ids": input_ids, "pixel_values": torch.ones(1, _HIDDEN_SIZE)} # [1,D] + if use_inputs_embeds: + kwargs["inputs_embeds"] = model.embedding(input_ids) # [1,3,D] + output = model(**kwargs) + expected = other_model(**kwargs) + assert torch.equal(output.logits, expected.logits) + output.logits.sum().backward() + for parameter in model.sound_und_model.parameters(): + assert parameter.grad is not None + assert torch.count_nonzero(parameter.grad) == 0 + assert model.sound_und_model.encoder.forward_calls == 1 + assert model.sound_und_model.projector.forward_calls == 1 + assert not model.embedding._forward_hooks + assert other_model.sound_und_model.encoder.forward_calls == 0 + assert other_model.sound_und_model.projector.forward_calls == 1 + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_trainable_audio_encoder_dummy_skipped_in_inference() -> None: + model = _FakeRawReasoner() + model.sound_und_model.encoder.requires_grad_(True) + patch_reasoner_audio_forward(model, audio_token_id=_AUDIO_TOKEN_ID, model_type="qwen3_vl") + with torch.no_grad(): + model(torch.tensor([[1, 2]])) + assert model.sound_und_model.encoder.forward_calls == 0 + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_trainable_audio_encoder_hook_removed_after_forward_failure() -> None: + model = _FakeRawReasoner() + model.sound_und_model.encoder.requires_grad_(True) + patch_reasoner_audio_forward(model, audio_token_id=_AUDIO_TOKEN_ID, model_type="qwen3_vl") + with pytest.raises(IndexError): + model(torch.tensor([[1000]])) + assert not model.embedding._forward_hooks diff --git a/cosmos_framework/model/generator/reasoner/qwen35_caption_accumulation_test.py b/cosmos_framework/model/generator/reasoner/qwen35_caption_accumulation_test.py index 3010af029..0342f0115 100644 --- a/cosmos_framework/model/generator/reasoner/qwen35_caption_accumulation_test.py +++ b/cosmos_framework/model/generator/reasoner/qwen35_caption_accumulation_test.py @@ -6,23 +6,28 @@ from contextlib import nullcontext from copy import deepcopy from types import SimpleNamespace -from typing import Any +from typing import Any, cast import pytest import torch import torch.nn as nn import torch.nn.functional as F +from torch.amp.grad_scaler import GradScaler from cosmos_framework.utils.config import CheckpointConfig from cosmos_framework.trainer import ImaginaireTrainer from cosmos_framework.utils.callback import Callback from cosmos_framework.model.generator.algorithm.loss import cross_entropy as ce -from cosmos_framework.configs.base.reasoner.defaults.policy_config import PolicyConfig, VLMModelConfig +from cosmos_framework.configs.base.reasoner.defaults.policy_config import ( + LBLConfig, + PolicyConfig, + VLMModelConfig, +) from cosmos_framework.model.generator.reasoner.qwen35_caption import Qwen35CaptionLoss from cosmos_framework.model.generator.vlm_model import VLMModel from cosmos_framework.utils.generator import flash_attn -pytestmark = [pytest.mark.L1, pytest.mark.CPU] +pytestmark = [pytest.mark.L0, pytest.mark.CPU] class _CaptionHead(nn.Module): @@ -56,7 +61,7 @@ def on_before_optimizer_step( model: VLMModel, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler, - grad_scaler: torch.amp.GradScaler, + grad_scaler: GradScaler, iteration: int = 0, ) -> None: self.gradients.append({name: parameter.grad.clone() for name, parameter in model.named_parameters()}) @@ -68,11 +73,13 @@ def _no_profile(*args: Any, **kwargs: Any) -> nullcontext[None]: return nullcontext() -def _make_model(monkeypatch: pytest.MonkeyPatch, exponent: float, normalize: bool = True) -> VLMModel: +def _make_model( + monkeypatch: pytest.MonkeyPatch, exponent: float, normalize: bool = True, *, model_type: str = "qwen3_5" +) -> VLMModel: def init_vlm(self: VLMModel, config: VLMModelConfig, checkpoint: CheckpointConfig) -> None: - self.model = _CaptionHead(exponent) - self.hf_config = SimpleNamespace(model_type="qwen3_5") - self.parallel_dims = None + setattr(self, "model", _CaptionHead(exponent)) + self.hf_config = SimpleNamespace(model_type=model_type) + setattr(self, "parallel_dims", None) def scalar_ce(hidden: torch.Tensor, weight: torch.Tensor, labels: torch.Tensor, ignore_index: int) -> torch.Tensor: return F.cross_entropy(F.linear(hidden, weight), labels, ignore_index=ignore_index, reduction="sum") @@ -90,7 +97,8 @@ def init_flash_attn_meta(deterministic: bool) -> None: enable_fused_weighted_ce=True, weighted_ce_exponent=exponent, normalize_weighted_ce_over_accumulation_window=normalize, - ) + ), + lbl=LBLConfig(coeff=0.1 if model_type == "qwen3_vl_moe" else None, method="global"), ), CheckpointConfig(), ) @@ -117,10 +125,10 @@ def test_accumulation_matches_full_batch_update( straggler_detection=SimpleNamespace(analyze_forward=False, analyze_backward=False, analyze_optimizer=False), ) ) - trainer.training_timer = _no_profile - trainer.straggler_detector = SimpleNamespace(profile_section=_no_profile) - trainer.callbacks = probe - scaler = torch.amp.GradScaler("cpu", enabled=False) + monkeypatch.setattr(trainer, "training_timer", _no_profile, raising=False) + monkeypatch.setattr(trainer, "straggler_detector", SimpleNamespace(profile_section=_no_profile), raising=False) + monkeypatch.setattr(trainer, "callbacks", probe, raising=False) + scaler = GradScaler("cpu", enabled=False) # Two consecutive windows also verify that the denominator is reset after each update. for window in range(2): @@ -132,8 +140,10 @@ def test_accumulation_matches_full_batch_update( for row, count in enumerate(counts): if count and not all_ignored: labels[row, -count:] = ids[row, -count:] - logits = reference.model.model(reference.model.embedding(ids)) + reference_head = cast(_CaptionHead, reference.model) + logits = reference_head.model(reference_head.embedding(ids)) reference_loss = ce.weighted_cross_entropy_loss(logits, labels, exponent) + assert isinstance(reference_loss, torch.Tensor) reference_loss.backward() expected_gradients = {name: parameter.grad.clone() for name, parameter in reference.named_parameters()} torch.nn.utils.clip_grad_norm_(reference.parameters(), max_norm=0.01) @@ -172,3 +182,37 @@ def test_window_normalization_can_be_disabled(monkeypatch: pytest.MonkeyPatch) - assert model._weighted_ce_window_denominator is None torch.testing.assert_close(output["train_objective_numerator"], loss.detach()) assert output["train_objective_denominator"].item() == 1 + + +@pytest.mark.parametrize("exponent", [0.0, 0.5, 1.0]) +def test_moe_accumulation_keeps_auxiliary_objective_separate(monkeypatch: pytest.MonkeyPatch, exponent: float) -> None: + torch.manual_seed(917) + model = _make_model(monkeypatch, exponent, model_type="qwen3_vl_moe") + reference = deepcopy(model) + assert model.requires_accumulation_window + + def auxiliary(self: VLMModel) -> torch.Tensor: + # A differentiable auxiliary objective isolates the trainer's CE/aux scaling + # from MoE routing kernels, whose distributed parity is tested separately. + return 0.1 * cast(_CaptionHead, self.model).model.weight.square().mean() + + monkeypatch.setattr(VLMModel, "_moe_load_balancing_loss", auxiliary) + ids = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]) + labels = torch.tensor([[-100, -100, -100, -100, 5], [-100, 7, 8, 9, 10]]) + reference_head = cast(_CaptionHead, reference.model) + reference_ce = ce.weighted_cross_entropy_loss(reference_head.model(reference_head.embedding(ids)), labels, exponent) + assert isinstance(reference_ce, torch.Tensor) + expected = reference_ce + auxiliary(reference) + expected.backward() + batches = tuple({"input_ids": ids[i : i + 1], "labels": labels[i : i + 1].clone()} for i in range(2)) + model.prepare_accumulation_window(batches) + for batch in batches: + output, loss = model.training_step(batch, iteration=0) + (output.get("_backward_loss", loss) / len(batches)).backward() + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0) + model.on_before_optimizer_step(optimizer, scheduler, iteration=0) + for actual, wanted in zip(model.parameters(), reference.parameters(), strict=True): + torch.testing.assert_close(actual.grad, wanted.grad, atol=2e-6, rtol=2e-5) + assert model._weighted_ce_window_plan is None + assert model._weighted_ce_window_denominator is None diff --git a/cosmos_framework/model/generator/reasoner/qwen3_omni_audio/configuration_qwen3_omni_audio.py b/cosmos_framework/model/generator/reasoner/qwen3_omni_audio/configuration_qwen3_omni_audio.py index c1c3888d4..c858b5d4c 100644 --- a/cosmos_framework/model/generator/reasoner/qwen3_omni_audio/configuration_qwen3_omni_audio.py +++ b/cosmos_framework/model/generator/reasoner/qwen3_omni_audio/configuration_qwen3_omni_audio.py @@ -15,6 +15,8 @@ Qwen3OmniMoeAudioEncoderConfig, ) +_ENCODER_OUTPUT_MODES = ("native", "pre_projection") + def get_qwen3_omni_thinking_audio_encoder_config() -> Qwen3OmniMoeAudioEncoderConfig: """Return the released Qwen3-Omni-30B-A3B-Thinking audio-tower config.""" @@ -56,6 +58,7 @@ def __init__( encoder_config: Qwen3OmniMoeAudioEncoderConfig | dict[str, Any] | None = None, projection_hidden_size: int = 4096, out_hidden_size: int = 2688, + encoder_output_mode: str = "native", **kwargs: Any, ) -> None: super().__init__(**kwargs) @@ -73,10 +76,19 @@ def __init__( _validate_positive_integer("encoder_config.output_dim", encoder_config.output_dim) _validate_positive_integer("projection_hidden_size", projection_hidden_size) _validate_positive_integer("out_hidden_size", out_hidden_size) - + if encoder_output_mode not in _ENCODER_OUTPUT_MODES: + raise ValueError(f"encoder_output_mode must be one of {_ENCODER_OUTPUT_MODES}, got {encoder_output_mode!r}") self.encoder_config = encoder_config self.projection_hidden_size = int(projection_hidden_size) self.out_hidden_size = int(out_hidden_size) + self.encoder_output_mode: str = encoder_output_mode + + @property + def encoder_output_size(self) -> int: + """Return the feature width emitted by the selected encoder output.""" + if self.encoder_output_mode == "pre_projection": + return int(self.encoder_config.d_model) + return int(self.encoder_config.output_dim) __all__ = ["Qwen3OmniAudioConfig", "get_qwen3_omni_thinking_audio_encoder_config"] diff --git a/cosmos_framework/model/generator/reasoner/qwen3_omni_audio/qwen3_omni_audio.py b/cosmos_framework/model/generator/reasoner/qwen3_omni_audio/qwen3_omni_audio.py index 1774d3c67..44da0e674 100644 --- a/cosmos_framework/model/generator/reasoner/qwen3_omni_audio/qwen3_omni_audio.py +++ b/cosmos_framework/model/generator/reasoner/qwen3_omni_audio/qwen3_omni_audio.py @@ -55,10 +55,14 @@ def __init__( self, config: Qwen3OmniMoeAudioEncoderConfig | None = None, encoder: Qwen3OmniMoeAudioEncoder | None = None, + output_mode: str = "native", ) -> None: super().__init__() + if output_mode not in ("native", "pre_projection"): + raise ValueError(f"Unsupported Qwen3-Omni encoder output mode: {output_mode!r}") self.config = config if config is not None else get_qwen3_omni_thinking_audio_encoder_config() self.encoder = encoder if encoder is not None else Qwen3OmniMoeAudioEncoder(self.config) + self.output_mode: str = output_mode def get_output_lengths(self, input_lengths: torch.Tensor) -> torch.Tensor: """Return the native encoded length for each input feature sequence.""" @@ -81,25 +85,43 @@ def forward( self, input_features: torch.Tensor, input_lengths: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: # input_features: [B,T,128], input_lengths: [B], returns [B,T',D], [B] """Encode padded log-mels and return padded native audio features.""" packed_features = torch.cat( [input_features[index, : int(length)].transpose(0, 1) for index, length in enumerate(input_lengths)], dim=1, - ).contiguous() - output_lengths = self.get_output_lengths(input_lengths) + ).contiguous() # [128,N] + output_lengths = self.get_output_lengths(input_lengths) # [B] use_bf16_autocast = ( self.config._attn_implementation == "flash_attention_2" and packed_features.device.type == "cuda" ) + pre_projection_features: list[torch.Tensor] = [] + + def capture_pre_projection(_module: nn.Module, inputs: tuple[torch.Tensor, ...]) -> None: # inputs[0]: [N,D] + pre_projection_features.append(inputs[0]) # [N,D] + + projection_hook = None + if self.output_mode == "pre_projection": + projection_hook = self.encoder.proj1.register_forward_pre_hook(capture_pre_projection) with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=use_bf16_autocast): - flat_features = self.encoder( - input_features=packed_features, - feature_lens=input_lengths, - ).last_hidden_state + try: + encoder_output = self.encoder( + input_features=packed_features, + feature_lens=input_lengths, + ) # last_hidden_state: [N,output_dim] + finally: + if projection_hook is not None: + projection_hook.remove() + if self.output_mode == "pre_projection": + if len(pre_projection_features) != 1: + raise RuntimeError("Qwen3-Omni native projection input was not captured exactly once") + flat_features = pre_projection_features[0] # [N,d_model] + else: + flat_features = encoder_output.last_hidden_state # [N,output_dim] padded_features = nn.utils.rnn.pad_sequence( - flat_features.split(output_lengths.tolist()), + list(flat_features.split(output_lengths.tolist())), batch_first=True, - ) + ) # [B,T',D] return padded_features, output_lengths @@ -108,7 +130,7 @@ class Qwen3OmniAudioProjector(AudioProjector): def __init__(self, config: Qwen3OmniAudioConfig) -> None: super().__init__( - input_hidden_size=config.encoder_config.output_dim, + input_hidden_size=config.encoder_output_size, projection_hidden_size=config.projection_hidden_size, out_hidden_size=config.out_hidden_size, ) @@ -125,7 +147,14 @@ def __init__( ) -> None: super().__init__() self.config = config if config is not None else Qwen3OmniAudioConfig() - self.encoder = encoder if encoder is not None else Qwen3OmniThinkingAudioEncoder(self.config.encoder_config) + self.encoder = ( + encoder + if encoder is not None + else Qwen3OmniThinkingAudioEncoder( + self.config.encoder_config, + output_mode=self.config.encoder_output_mode, + ) + ) self.projector = projector if projector is not None else Qwen3OmniAudioProjector(self.config) def init_weights(self, buffer_device: torch.device | None = None) -> None: @@ -133,6 +162,13 @@ def init_weights(self, buffer_device: torch.device | None = None) -> None: self.encoder.reset_nonpersistent_buffers(buffer_device) self.projector.reset_parameters() + def configure_encoder_trainability(self, *, train_encoder: bool) -> None: + """Train the effective trunk while keeping bypassed native projections frozen.""" + self.encoder.requires_grad_(train_encoder) + if self.config.encoder_output_mode == "pre_projection": + self.encoder.encoder.proj1.requires_grad_(False) + self.encoder.encoder.proj2.requires_grad_(False) + def forward( self, audio_features: torch.Tensor, diff --git a/cosmos_framework/model/generator/reasoner/qwen3_vl_fused_ce.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_fused_ce.py new file mode 100644 index 000000000..8510f7c8d --- /dev/null +++ b/cosmos_framework/model/generator/reasoner/qwen3_vl_fused_ce.py @@ -0,0 +1,64 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Opt-in fused SFT objective for dense and MoE Qwen3-VL models.""" + +from functools import wraps +from types import MethodType +from typing import Any, cast + +import torch.nn as nn +from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration +from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import Qwen3VLMoeForConditionalGeneration + +from cosmos_framework.model.generator.algorithm.loss.cross_entropy import fused_weighted_cross_entropy_loss +from cosmos_framework.model.generator.reasoner.qwen3_vl_rope import configure_qwen3_vl_fp32_rope +from cosmos_framework.model.generator.reasoner.qwen35_caption import Qwen35CaptionLoss + + +def configure_qwen3_vl_fused_ce( + model: nn.Module, *, exponent: float, recompute: bool = False, accumulate_head_grad_fp32: bool = True +) -> None: + """Evaluate LM projection and CE while the FSDP root's weight is gathered. + + The existing Qwen3-VL backbone patch owns text-only vision dependencies. + MoE routing hooks remain on the original decoder blocks, so VLMModel still + computes its configured load-balancing objective outside this wrapper. + """ + if getattr(model, "config").model_type not in {"qwen3_vl", "qwen3_vl_moe"}: + raise ValueError("Qwen3-VL fused CE requires a dense or MoE Qwen3-VL model") + if getattr(model, "_qwen3_vl_fused_ce", False): + raise ValueError("Qwen3-VL fused CE is already configured") + if not 0.0 <= exponent <= 1.0: + raise ValueError("Weighted CE exponent must be in [0,1]") + # The family check establishes the native Qwen3-VL model interface. + model = cast(Qwen3VLForConditionalGeneration | Qwen3VLMoeForConditionalGeneration, model) + # Evaluate native position phases in FP32 on the fused path, using the same + # phases in CP=1 and CP>1 when comparing the fused path. + configure_qwen3_vl_fp32_rope(model.model.language_model.rotary_emb) + original_forward = model.forward + + @wraps(original_forward.__func__) + def forward( + self: Qwen3VLForConditionalGeneration | Qwen3VLMoeForConditionalGeneration, *args: Any, **kwargs: Any + ) -> Qwen35CaptionLoss: + labels = kwargs.pop("labels", None) + if labels is None: + raise ValueError("Fused Qwen3-VL SFT requires labels") + kwargs.pop("logits_to_keep", None) + kwargs["use_cache"] = False + kwargs["return_dict"] = True + hidden = self.model(*args, **kwargs).last_hidden_state + loss, stats = fused_weighted_cross_entropy_loss( + hidden, + labels, + self.lm_head.weight, + exponent=exponent, + context_parallel_group=getattr(self, "_cosmos_cp_group", None), + cu_seq_lens=kwargs.get("cu_seq_lens_q"), + recompute=recompute, + accumulate_head_grad_fp32=accumulate_head_grad_fp32 and kwargs.get("cu_seq_lens_q") is not None, + ) + return Qwen35CaptionLoss(loss, stats) + + model.forward = MethodType(forward, model) + setattr(model, "_qwen3_vl_fused_ce", True) diff --git a/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/fixed_sum.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/fixed_sum.py new file mode 100644 index 000000000..a50c07930 --- /dev/null +++ b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/fixed_sum.py @@ -0,0 +1,190 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Opt-in fixed-order FP32 dispatch/combine for grouped Qwen3-VL MoE experts. + +The expert GEMMs and parameters are unchanged. Only top-k accumulation order +and accumulation precision change; this does not promise deterministic attention. + +A token is routed to K experts. Combining their outputs, and adding the K input +gradients in dispatch backward, each requires a sum over those assignments. +This module gathers them in top-k order and accumulates in FP32 instead of using +unordered low-precision scatter additions. It is an optional Qwen MoE numerical +control, useful for CP comparisons but equally applicable to CP1. CP's shared +sequence/head communication does not perform this expert reduction. + +``assignments[row] = token * K + slot`` maps an expert-grouped, padded row back +to a top-k assignment. ``tokens * K`` denotes a padding row. ``inverse`` maps +each real assignment back to its grouped row, or -1 for a masked assignment. +""" + +from collections.abc import Callable +from typing import Any, cast + +import torch +import triton +import triton.language as tl +from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock + +from cosmos_framework.model.generator.moe.native import NativeMoEExperts +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe import ( + _run_experts_grouped_mm, +) +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe_kernels import ( + TOKEN_GROUP_ALIGN_SIZE_M, + _generate_permute_indices, +) + + +@triton.jit +def _inverse_kernel( + assignments: torch.Tensor, inverse: torch.Tensor, N: tl.constexpr, P: tl.constexpr, BLOCK: tl.constexpr +) -> None: + """Invert real assignment IDs; padded rows must not overwrite the inverse.""" + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + assignment = tl.load(assignments + row, row < P, other=N) + tl.store(inverse + assignment, row, (row < P) & (assignment < N)) + + +@triton.jit +def _combine_kernel( + values: torch.Tensor, + inverse: torch.Tensor, + result: torch.Tensor, + N: tl.constexpr, + D: tl.constexpr, + K: tl.constexpr, + BLOCK: tl.constexpr, +) -> None: + """Sum each token's K rows in slot order and FP32, then cast on output store.""" + offset = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + token, column = offset // D, offset % D + total = tl.full((BLOCK,), 0, tl.float32) + for k in range(K): + row = tl.load(inverse + token * K + k, token < N, other=-1) + value = tl.load(values + row * D + column, (token < N) & (row >= 0), other=0).to(tl.float32) + total = total + value + tl.store(result + offset, total, token < N) + + +def inverse_assignments(assignments: torch.Tensor, tokens: int, top_k: int) -> torch.Tensor: + """Build the [tokens * top_k] lookup, leaving absent/masked assignments at -1.""" + inverse = torch.full((tokens * top_k,), -1, dtype=torch.int64, device=assignments.device) + # Triton specializes Python integer launch arguments annotated tl.constexpr in the kernel. + launch = cast(Callable[..., None], _inverse_kernel[(triton.cdiv(assignments.numel(), 256),)]) + launch(assignments, inverse, tokens * top_k, assignments.numel(), 256) + return inverse + + +def fixed_sum(values: torch.Tensor, inverse: torch.Tensor, tokens: int, top_k: int) -> torch.Tensor: + """Reduce grouped [padded_rows, hidden] values to [tokens, hidden] in their input dtype. + + Only accumulation is FP32; absent assignments contribute zero. Autograd is + supplied by Dispatch/Combine because this Triton call has no backward itself. + """ + result = values.new_empty((tokens, values.shape[1])) + launch = cast(Callable[..., None], _combine_kernel[(triton.cdiv(result.numel(), 256),)]) + launch(values.contiguous(), inverse, result, tokens, values.shape[1], top_k, 256) + return result + + +class Dispatch(torch.autograd.Function): + """Gather token inputs into expert order; backward sums their K gradients. + + A sentinel zero token supplies alignment padding without routing a real token. + """ + + @staticmethod + def forward( + ctx: Any, + hidden: torch.Tensor, + assignments: torch.Tensor, + inverse: torch.Tensor, + top_k: int, + ) -> torch.Tensor: + ctx.save_for_backward(assignments, inverse) + ctx.tokens, ctx.top_k = hidden.shape[0], top_k + padded = torch.cat([hidden, hidden.new_zeros((1, hidden.shape[1]))]) + return padded.index_select(0, assignments // top_k) + + @staticmethod + def backward(ctx: Any, grad: torch.Tensor) -> tuple[torch.Tensor, None, None, None]: + _, inverse = ctx.saved_tensors + return fixed_sum(grad, inverse, ctx.tokens, ctx.top_k), None, None, None + + +class Combine(torch.autograd.Function): + """Sum expert outputs per token; backward copies its gradient to each assignment. + + Routing scores have already been applied by the grouped expert computation. + """ + + @staticmethod + def forward( + ctx: Any, + values: torch.Tensor, + assignments: torch.Tensor, + inverse: torch.Tensor, + tokens: int, + top_k: int, + ) -> torch.Tensor: + ctx.save_for_backward(assignments) + ctx.top_k = top_k + return fixed_sum(values, inverse, tokens, top_k) + + @staticmethod + def backward(ctx: Any, grad: torch.Tensor) -> tuple[torch.Tensor, None, None, None, None]: + (assignments,) = ctx.saved_tensors + padded = torch.cat([grad, grad.new_zeros((1, grad.shape[1]))]) + return padded.index_select(0, assignments // ctx.top_k), None, None, None, None + + +def deterministic_experts_forward( + self: NativeMoEExperts, + hidden_states: torch.Tensor, + topk_scores: torch.Tensor, + expert_indices: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + token_mask: torch.Tensor | None, +) -> torch.Tensor: + """Reuse native grouped expert GEMMs with fixed-order dispatch/combine reductions. + + The stable expert ordering and inverse lookup preserve top-k slot identity + through alignment padding. Masked tokens have no real expert assignments. + Expert weights, routing scores and the auxiliary objective are unchanged; + the name refers to the reductions, not whole-model determinism. + """ + tokens, _ = hidden_states.shape + top_k = self.top_k + if token_mask is not None: + expert_indices = expert_indices.masked_fill(~token_mask.unsqueeze(1), self.num_experts) + order = torch.argsort(expert_indices.reshape(-1), stable=True) + sorted_scores = topk_scores.reshape(-1)[order] + alignment = TOKEN_GROUP_ALIGN_SIZE_M + experts = num_tokens_per_expert.shape[0] + padded_size = ((tokens * top_k + experts * alignment + alignment - 1) // alignment) * alignment + permutation, counts = _generate_permute_indices(num_tokens_per_expert, experts, padded_size, alignment) + order_ext = torch.cat([order, order.new_full((1,), tokens * top_k)]) + assignments = order_ext[permutation.long()] + inverse = inverse_assignments(assignments, tokens, top_k) + scores_ext = torch.cat([sorted_scores, sorted_scores.new_zeros(1)]) + scores = scores_ext[permutation.long()] + routed = Dispatch.apply(hidden_states, assignments, inverse, top_k) + output = _run_experts_grouped_mm(self.gate_up_proj, self.down_proj, self.act_fn, routed, counts, scores) + return Combine.apply(output, assignments, inverse, tokens, top_k) + + +def configure_qwen3_vl_moe_fixed_sum(model: torch.nn.Module) -> int: + """Enable the optional expert reduction after grouped-mm setup, before AC/FSDP. + + Only this model instance changes. The default recipe does not opt in, and no + parameter or checkpoint layout changes when the option is enabled. + """ + blocks = [module for module in model.modules() if isinstance(module, Qwen3VLMoeTextSparseMoeBlock)] + if not blocks: + raise ValueError("Fixed MoE summation requires Qwen3-VL-MoE blocks") + if any(not hasattr(block.experts, "_reorder_tokens") for block in blocks): + raise ValueError("Configure grouped-mm experts before fixed MoE summation") + for block in blocks: + experts = block.experts + experts.forward = deterministic_experts_forward.__get__(experts, type(experts)) + return len(blocks) diff --git a/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_test.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_test.py index 5b455c98b..b844df021 100644 --- a/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_test.py +++ b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_test.py @@ -21,6 +21,7 @@ Qwen3VLMoeTextSparseMoeBlock, _weighted_expert_counts, ) +from cosmos_framework.model.generator.utils.moe_utils import preserve_moe_learning_accumulators def test_router_activation_defaults_to_softmax() -> None: @@ -106,6 +107,92 @@ def test_aux_loss_free_controller_uses_block_config() -> None: assert "tokens_per_expert" not in block.state_dict() +def _make_router_learning_block() -> Qwen3VLMoeTextSparseMoeBlock: + model_config = Qwen3VLMoeTextConfig( + hidden_size=4, + moe_intermediate_size=4, + num_experts=4, + num_experts_per_tok=2, + hidden_act="silu", + ) + return Qwen3VLMoeTextSparseMoeBlock( + model_config, + cosine_router_config=CosineRouterConfig(input_centering="ema"), + aux_loss_free_load_balancing_config=AuxLossFreeLoadBalancingConfig( + enabled=True, + update_speed=0.25, + max_bias=None, + ), + ) + + +def _seed_router_learning_accumulators(block: Qwen3VLMoeTextSparseMoeBlock) -> None: + block.tokens_per_expert.copy_(torch.tensor([3.0, 1.0, 0.0, 0.0])) # [E] + block.cosine_router.router_bias_sum.copy_(torch.tensor([2.0, 4.0, 6.0, 8.0])) # [D] + block.cosine_router.router_bias_count.copy_(torch.tensor([2.0])) # [1] + + +@pytest.mark.L0 +def test_sampling_observation_preserves_router_updates_but_keeps_monitoring() -> None: + observed_block = _make_router_learning_block() + baseline_block = _make_router_learning_block() + observed_block.experts = _RecordingExperts() + _seed_router_learning_accumulators(observed_block) + _seed_router_learning_accumulators(baseline_block) + observed_net = nn.Module() + observed_net.add_module("moe_gen", observed_block) + observation_hidden_states = torch.arange(20, dtype=torch.float32).reshape(5, 4) # [N,D] + + with preserve_moe_learning_accumulators(observed_net): + observation_output, _ = observed_block(observation_hidden_states) # [N,D], LBLMetadata + + assert observation_output.shape == observation_hidden_states.shape + + torch.testing.assert_close(observed_block.tokens_per_expert, baseline_block.tokens_per_expert) + torch.testing.assert_close( + observed_block.cosine_router.router_bias_sum, + baseline_block.cosine_router.router_bias_sum, + ) + torch.testing.assert_close( + observed_block.cosine_router.router_bias_count, + baseline_block.cosine_router.router_bias_count, + ) + # Observation telemetry is deliberately not part of the restored learning state. + assert observed_block.total_tokens_per_expert.sum().item() == 10 + torch.testing.assert_close(observed_block.total_tokens, torch.tensor([5])) # [1] + + observed_block.update_bias() + baseline_block.update_bias() + observed_block.update_router_bias(device_mesh=None) + baseline_block.update_router_bias(device_mesh=None) + + torch.testing.assert_close(observed_block.expert_bias, baseline_block.expert_bias) + torch.testing.assert_close(observed_block.cosine_router.router_bias, baseline_block.cosine_router.router_bias) + + +@pytest.mark.L0 +def test_sampling_observation_restores_learning_accumulators_after_exception() -> None: + block = _make_router_learning_block() + block.experts = _RecordingExperts() + _seed_router_learning_accumulators(block) + net = nn.Module() + net.add_module("moe_gen", block) + expected_counts = block.tokens_per_expert.clone() # [E] + expected_sum = block.cosine_router.router_bias_sum.clone() # [D] + expected_count = block.cosine_router.router_bias_count.clone() # [1] + + with pytest.raises(RuntimeError, match="sampling failed"): + with preserve_moe_learning_accumulators(net): + hidden_states = torch.randn(3, 4) # [N,D] + block(hidden_states) + raise RuntimeError("sampling failed") + + torch.testing.assert_close(block.tokens_per_expert, expected_counts) + torch.testing.assert_close(block.cosine_router.router_bias_sum, expected_sum) + torch.testing.assert_close(block.cosine_router.router_bias_count, expected_count) + assert block.training is True + + # ----------------------------------------------------------------------------- # Padded streams # diff --git a/cosmos_framework/model/generator/reasoner/qwen3_vl_rope.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_rope.py new file mode 100644 index 000000000..afa37532b --- /dev/null +++ b/cosmos_framework/model/generator/reasoner/qwen3_vl_rope.py @@ -0,0 +1,41 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Shape-independent FP32 phases for Qwen3-VL long-context RoPE.""" + +from types import MethodType +from typing import cast + +import torch +import torch.nn as nn +from transformers.modeling_rope_utils import dynamic_rope_update +from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextRotaryEmbedding +from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import Qwen3VLMoeTextRotaryEmbedding + + +@torch.no_grad() +@dynamic_rope_update +def _fp32_rotary_forward( + self: Qwen3VLTextRotaryEmbedding | Qwen3VLMoeTextRotaryEmbedding, x: torch.Tensor, position_ids: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + # Qwen3-VL uses separate temporal, height and width position axes. + if position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + device_type = x.device.type if x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + # The native outer-product matmul may use TF32 despite autocast being off. + # Long position IDs then lose precision differently for full sequences and + # CP shards. Elementwise FP32 multiplication avoids TF32 without changing + # the precision or performance policy for the model's other matmuls. + # Native Qwen rotary modules register inv_freq as a tensor buffer. + inv_freq = cast(torch.Tensor, self.inv_freq) + frequencies = position_ids.float().unsqueeze(-1) * inv_freq.float()[None, None, None, :] + frequencies = self.apply_interleaved_mrope(frequencies, self.mrope_section) + phases = torch.cat((frequencies, frequencies), dim=-1) + cos = phases.cos() * self.attention_scaling + sin = phases.sin() * self.attention_scaling + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +def configure_qwen3_vl_fp32_rope(rotary_embedding: nn.Module) -> None: + """Keep native MRoPE and dynamic updates, replacing only the phase outer product.""" + rotary_embedding.forward = MethodType(_fp32_rotary_forward, rotary_embedding) diff --git a/cosmos_framework/model/generator/tokenizers/lidar/lidar_tokenizer_v1.py b/cosmos_framework/model/generator/tokenizers/lidar/lidar_tokenizer_v1.py index f72edcd81..b6d059eb4 100644 --- a/cosmos_framework/model/generator/tokenizers/lidar/lidar_tokenizer_v1.py +++ b/cosmos_framework/model/generator/tokenizers/lidar/lidar_tokenizer_v1.py @@ -143,7 +143,8 @@ def __init__( super().__init__(object_store_credential_path_pretrained) self.device = torch.device(device) - self.dtype = as_torch_dtype(dtype) + # Exported configs serialize torch dtypes as strings (for example, "float32"). + self.dtype: torch.dtype = as_torch_dtype(dtype) self.sample_posterior = sample_posterior self.apply_validity_mask = apply_validity_mask if streaming_chunk_frames is not None and streaming_chunk_frames < 1: diff --git a/cosmos_framework/model/generator/utils/component_unipc.py b/cosmos_framework/model/generator/utils/component_unipc.py new file mode 100644 index 000000000..95c64ae0d --- /dev/null +++ b/cosmos_framework/model/generator/utils/component_unipc.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Inference coordination for jointly evaluated components with separate UniPC schedules.""" + +import math +from collections.abc import Callable +from dataclasses import dataclass + +import torch + +from cosmos_framework.utils.progress_bar import progress_bar +from cosmos_framework.model.generator.diffusion.samplers.fm_solvers_unipc import FlowUniPCMultistepScheduler +from cosmos_framework.model.generator.diffusion.samplers.unipc import UniPCSamplerConfig + + +@dataclass(frozen=True) +class ComponentSchedule: + """Shared component shifts and the corresponding sizes in each flat sample.""" + + shifts: tuple[float, ...] + lengths: list[tuple[int, ...]] + + +@torch.no_grad() +def sample_component_unipc( + velocity_fn: Callable[[list[torch.Tensor], list[torch.Tensor]], list[torch.Tensor]], + noise: list[torch.Tensor], # B x [N_i] + *, + cfg: UniPCSamplerConfig, + schedule: ComponentSchedule, + num_steps: int, + step_callback: Callable[[int, int], None] | None = None, +) -> list[torch.Tensor]: # B x [N_i] + """Advance generic flat components with one joint velocity evaluation per step. + + The caller defines component boundaries and adapts the component-ordered + timesteps (each [1,1]) to its velocity function. Each sample/component owns + an unchanged UniPC scheduler, including its sigma grid and multistep history. + No common-clock derivative approximation or additional model call is used. + """ + if num_steps <= 0: + raise ValueError("num_steps must be positive") + if cfg.use_dynamic_shifting: + raise ValueError("Component UniPC schedules require fixed shifts") + if not schedule.shifts or any(not math.isfinite(value) or value <= 0 for value in schedule.shifts): + raise ValueError("Component shifts must be positive and finite") + if not noise or len(noise) != len(schedule.lengths): + raise ValueError("Component sampler sample metadata lengths differ") + + schedulers: list[list[FlowUniPCMultistepScheduler]] = [] + for sample, lengths in zip(noise, schedule.lengths, strict=True): + if ( + sample.ndim != 1 + or len(lengths) != len(schedule.shifts) + or any(length <= 0 for length in lengths) + or sum(lengths) != sample.numel() + ): + raise ValueError("Component lengths must partition each flat sample into nonempty components") + sample_schedulers: list[FlowUniPCMultistepScheduler] = [] + for shift in schedule.shifts: + scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=cfg.num_train_timesteps, + shift=cfg.shift, + use_dynamic_shifting=False, + ) + scheduler.set_timesteps(num_steps, device=sample.device, shift=shift) + sample_schedulers.append(scheduler) + schedulers.append(sample_schedulers) + + latent = noise # B x [N_i] + for step_index in progress_bar(range(num_steps), desc="Sampling", total=num_steps): + if step_callback is not None: + step_callback(step_index, num_steps) + timesteps = [scheduler.timesteps[step_index].reshape(1, 1) for scheduler in schedulers[0]] # K x [1,1] + predictions = velocity_fn(latent, timesteps) # B x [N_i] + if len(predictions) != len(latent): + raise ValueError("Component sampler prediction batch length differs") + updated: list[torch.Tensor] = [] + for sample, prediction, lengths, sample_schedulers in zip( + latent, predictions, schedule.lengths, schedulers, strict=True + ): + if prediction.shape != sample.shape: + raise ValueError("Component sampler prediction shape differs from its sample") + sample_parts = sample.split(lengths) # K x [N_ik] + prediction_parts = prediction.split(lengths) # K x [N_ik] + parts: list[torch.Tensor] = [] + for scheduler, part, velocity in zip(sample_schedulers, sample_parts, prediction_parts, strict=True): + # Native UniPC's multistep einsums require a leading batch dimension. + parts.append( + scheduler.step( + model_output=velocity, + timestep=scheduler.timesteps[step_index], + sample=part.unsqueeze(0), # [1,N_ik] + return_dict=False, + )[0].squeeze(0) # [N_ik] + ) + updated.append(torch.cat(parts)) # [N_i] + latent = updated # B x [N_i] + return latent diff --git a/cosmos_framework/model/generator/utils/component_unipc_test.py b/cosmos_framework/model/generator/utils/component_unipc_test.py new file mode 100644 index 000000000..608002c44 --- /dev/null +++ b/cosmos_framework/model/generator/utils/component_unipc_test.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +import pytest +import torch + +from cosmos_framework.model.generator.diffusion.samplers.unipc import UniPCSampler, UniPCSamplerConfig +from cosmos_framework.model.generator.utils.component_unipc import ComponentSchedule, sample_component_unipc + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + + +@pytest.mark.parametrize("shift", [1.0, 5.0]) +@pytest.mark.parametrize("num_steps", [1, 7]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_equal_shifts_match_legacy_unipc(shift: float, num_steps: int, dtype: torch.dtype) -> None: + noise = [torch.linspace(-1, 1, 8, dtype=dtype), torch.linspace(-2, 2, 11, dtype=dtype)] # [8], [11] + sampler = UniPCSampler(tensor_kwargs={"device": torch.device("cpu")}) + + def velocity(values: list[torch.Tensor], timestep: torch.Tensor) -> list[torch.Tensor]: # B x [N_i], [1,1] + # Every component depends on all the others, so this checks joint evolution. + return [0.2 * value.sin() + 0.1 * value.mean() + timestep.item() / 5000 for value in values] # B x [N_i] + + def joint_velocity( + values: list[torch.Tensor], + timesteps: list[torch.Tensor], # B x [N_i], K x [1,1] + ) -> list[torch.Tensor]: # B x [N_i] + torch.testing.assert_close(timesteps[0], timesteps[1]) + return velocity(values, timesteps[0]) + + expected = sampler(velocity, noise, num_steps=num_steps, shift=shift, seed=[3, 7]) # B x [N_i] + actual = sample_component_unipc( + joint_velocity, + noise, + cfg=sampler.cfg, + schedule=ComponentSchedule(lengths=[(3, 5), (7, 4)], shifts=(shift, shift)), + num_steps=num_steps, + ) # B x [N_i] + for result, reference in zip(actual, expected, strict=True): + torch.testing.assert_close(result, reference, rtol=0, atol=0) + + +def test_unequal_shifts_match_individual_samplers_with_one_joint_call_per_step() -> None: + lengths = [(2, 3, 4), (5, 2, 3)] + shifts = (1.0, 3.0, 5.0) + seeds = [3, 7] + noise = [torch.linspace(-1, 1, sum(parts), dtype=torch.float64) for parts in lengths] # B x [N_i] + sampler = UniPCSampler(tensor_kwargs={"device": torch.device("cpu")}) + callbacks: list[tuple[int, int]] = [] + model_timesteps: list[tuple[int, ...]] = [] + + def velocity(value: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: # [N_ik], [1,1] -> [N_ik] + return 0.2 * value.sin() + timestep.item() / 5000 # [N_ik] + + def joint_velocity( + values: list[torch.Tensor], + timesteps: list[torch.Tensor], # B x [N_i], K x [1,1] + ) -> list[torch.Tensor]: # B x [N_i] + assert callbacks[-1] == (len(model_timesteps), 7) + model_timesteps.append(tuple(int(t.item()) for t in timesteps)) + return [ + torch.cat([velocity(part, t) for part, t in zip(value.split(sizes), timesteps, strict=True)]) + for value, sizes in zip(values, lengths, strict=True) + ] # B x [N_i] + + actual = sample_component_unipc( + joint_velocity, + noise, + cfg=sampler.cfg, + schedule=ComponentSchedule(lengths=lengths, shifts=shifts), + num_steps=7, + step_callback=lambda step, count: callbacks.append((step, count)), + ) # B x [N_i] + assert callbacks == [(step, 7) for step in range(7)] + assert len(model_timesteps) == 7 + assert any(len(set(timesteps)) == 3 for timesteps in model_timesteps) + for sample, result, sizes, seed in zip(noise, actual, lengths, seeds, strict=True): + for source, output, shift in zip(sample.split(sizes), result.split(sizes), shifts, strict=True): + expected = sampler(velocity, source, num_steps=7, shift=shift, seed=seed) # [N_ik] + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("shifts", [(0.0,), (-1.0,), (float("inf"),), (float("nan"),), ()]) +def test_invalid_shifts_fail_before_velocity_evaluation(shifts: tuple[float, ...]) -> None: + def velocity(values: list[torch.Tensor], timesteps: list[torch.Tensor]) -> list[torch.Tensor]: + pytest.fail("Invalid schedules must fail before evaluating the model") + + with pytest.raises(ValueError, match="positive and finite"): + sample_component_unipc( + velocity, + [torch.zeros(3)], # 1 x [3] + cfg=UniPCSamplerConfig(), + schedule=ComponentSchedule(lengths=[(3,)], shifts=shifts), + num_steps=3, + ) diff --git a/cosmos_framework/model/generator/utils/kv_cache.py b/cosmos_framework/model/generator/utils/kv_cache.py index fd5cf4ad5..bb893d4a7 100644 --- a/cosmos_framework/model/generator/utils/kv_cache.py +++ b/cosmos_framework/model/generator/utils/kv_cache.py @@ -21,7 +21,6 @@ import torch.nn.functional as F from cosmos_framework.configs.base.defaults.replay_attention import TeacherForcingReplayPolicyConfig -from cosmos_framework.data.generator.sequence_packing.runtime import get_num_real_samples, to_device_nonblocking from cosmos_framework.model.generator.utils.kv_storage_backend import ( BF16StorageBackend, FP8StorageBackend, @@ -34,6 +33,7 @@ RollingPromptLayout, RollingTextSinkLayer, ) +from cosmos_framework.data.generator.sequence_packing.runtime import get_num_real_samples, to_device_nonblocking # See cosmos_framework/data/generator/sequence_packing/ for the canonical definition. SequencePack = dict[str, Any] diff --git a/cosmos_framework/model/generator/utils/kv_cache_test.py b/cosmos_framework/model/generator/utils/kv_cache_test.py index 220554983..916276413 100644 --- a/cosmos_framework/model/generator/utils/kv_cache_test.py +++ b/cosmos_framework/model/generator/utils/kv_cache_test.py @@ -8,15 +8,8 @@ import pytest import torch -from cosmos_framework.configs.base.defaults.replay_attention import TeacherForcingReplayPolicyConfig -from cosmos_framework.data.generator.sequence_packing.runtime import ( - SequencePack, - get_gen_seq, - get_num_real_samples, - has_pad_segment, - sequence_pack_from_packed_sequence, -) from cosmos_framework.model.attention import attention +from cosmos_framework.configs.base.defaults.replay_attention import TeacherForcingReplayPolicyConfig from cosmos_framework.model.generator.mot.attention import SplitInfo from cosmos_framework.model.generator.mot.causal_attention import ( attention_AR_gen_only, @@ -44,6 +37,13 @@ BF16StorageBackend, FP8StorageBackend, ) +from cosmos_framework.data.generator.sequence_packing.runtime import ( + SequencePack, + get_gen_seq, + get_num_real_samples, + has_pad_segment, + sequence_pack_from_packed_sequence, +) @pytest.mark.L0 diff --git a/cosmos_framework/model/generator/utils/memory.py b/cosmos_framework/model/generator/utils/memory.py index 115a8d332..f206a0fb7 100644 --- a/cosmos_framework/model/generator/utils/memory.py +++ b/cosmos_framework/model/generator/utils/memory.py @@ -20,7 +20,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from dataclasses import dataclass +from dataclasses import dataclass, field import torch @@ -29,6 +29,30 @@ KVToStore = tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] +@dataclass +class ARTextLayerCache: + """Invariant caption results for one noisy chunk, CFG branch and decoder layer. + + The owning rolling request is discarded at the chunk boundary. Clean forwards + have no text cache. Sequence-sharded projections retain native RoPE and norms; + exchanged tensors retain native CP KV-head replication and text-sink behavior. + No media activations or historical media K/V are stored here. + The first noisy forward populates this state; subsequent steps read it. + """ + + projected: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = ( + None # Q: [S_text_local,H,D]; K,V: [S_text_local,H_kv,D] + ) + projected_normalized_k: torch.Tensor | None = None # [S_text_local,H_kv,D] + exchanged: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = ( + None # Q: [S_text,H_local,D]; K,V: [S_text,H_kv_local,D] + ) + exchanged_normalized_k: torch.Tensor | None = None # [S_text,H_kv_local,D] + attention_output: torch.Tensor | None = None # [S_text,H_local*D] + scattered_output: torch.Tensor | None = None # [S_text_local,H,D] + hidden_output: torch.Tensor | None = None # [S_text_local,hidden_size] + + @dataclass class MemoryValue(ABC): """Read-only tensor container safe to pass into ``torch.compile``. @@ -38,6 +62,10 @@ class MemoryValue(ABC): mutate state or perform non-trivial computation. """ + # Inference may attach request-owned caption state. Layers populate it once + # per chunk, including under compile; media-memory tensors remain read-only. + text_cache: ARTextLayerCache | None = field(default=None, kw_only=True) + @property def supports_context_parallel_attention(self) -> bool: """Whether this memory value is compatible with context-parallel attention. diff --git a/cosmos_framework/model/generator/utils/model_construction.py b/cosmos_framework/model/generator/utils/model_construction.py new file mode 100644 index 000000000..29e83926c --- /dev/null +++ b/cosmos_framework/model/generator/utils/model_construction.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Construction scope for inspecting model architecture during checkpoint conversion.""" + +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar + +_architecture_only: ContextVar[bool] = ContextVar("architecture_only_model_construction", default=False) + + +def is_architecture_only_model_construction() -> bool: + """Whether construction may skip execution-backend availability checks.""" + return _architecture_only.get() + + +@contextmanager +def architecture_only_model_construction() -> Iterator[None]: + """Skip runtime backend checks for architecture inspection, without changing model config. + + Callers must separately arrange empty-weight construction and must not execute + the resulting model. Normal training and inference must keep backend validation. + """ + token = _architecture_only.set(True) + try: + yield + finally: + _architecture_only.reset(token) diff --git a/cosmos_framework/model/generator/utils/moe_utils.py b/cosmos_framework/model/generator/utils/moe_utils.py index 94f428bcc..82659653c 100644 --- a/cosmos_framework/model/generator/utils/moe_utils.py +++ b/cosmos_framework/model/generator/utils/moe_utils.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: OpenMDW-1.1 from collections.abc import Iterator +from contextlib import contextmanager import torch from torch.distributed.device_mesh import DeviceMesh @@ -24,6 +25,40 @@ def _iter_generation_moe_blocks( yield name, module +@contextmanager +def preserve_moe_learning_accumulators(net: torch.nn.Module) -> Iterator[None]: + """Restore router learning accumulators after an observation-only forward. + + Sampling callbacks intentionally run with the model's existing module modes and RNG + behavior, so changing the model to evaluation mode would also change their outputs. + Instead, preserve only the transient state consumed by the next optimizer-step router + updates. Monitoring buffers remain live so observations can still be inspected. + """ + snapshots: list[tuple[torch.Tensor, torch.Tensor]] = [] + for _, module in _iter_generation_moe_blocks(net): + if module.aux_loss_free_load_balancing_config.enabled: + snapshots.append((module.tokens_per_expert, module.tokens_per_expert.detach().clone())) # [E], [E] + if module.cosine_router.input_centering == "ema": + snapshots.extend( + [ + ( + module.cosine_router.router_bias_sum, + module.cosine_router.router_bias_sum.detach().clone(), + ), # [D], [D] + ( + module.cosine_router.router_bias_count, + module.cosine_router.router_bias_count.detach().clone(), + ), # [1], [1] + ] + ) + try: + yield + finally: + with torch.no_grad(): + for accumulator, snapshot in snapshots: + accumulator.copy_(snapshot) # [*accumulator_shape] + + def set_hf_moe_token_mask(net: torch.nn.Module, attention_mask: torch.Tensor | None) -> None: """Publish which rows are real tokens to every patched HF sparse MoE block. diff --git a/cosmos_framework/model/generator/utils/multiview_ar.py b/cosmos_framework/model/generator/utils/multiview_ar.py index ad9c37d56..5625270bd 100644 --- a/cosmos_framework/model/generator/utils/multiview_ar.py +++ b/cosmos_framework/model/generator/utils/multiview_ar.py @@ -1955,6 +1955,11 @@ def iter_samples_multiview_transfer_autoregressive( has_negative_prompt, caption_groups, ) + # Optional benchmark boundary: input encoding is done; clean-prefix work + # and diffusion have not started. Ordinary inference installs no callback. + on_inputs_prepared: Callable[[], None] | None = data_batch.pop("_on_ar_inputs_prepared", None) + if on_inputs_prepared is not None: + on_inputs_prepared() cfgp_enabled = host.parallel_dims is not None and host.parallel_dims.cfgp_enabled if cfgp_enabled: seed = _broadcast_seed([seed], host.parallel_dims.cfgp_mesh.get_group(), host.parallel_dims.cfgp_rank)[0] diff --git a/cosmos_framework/model/generator/utils/rolling_kv/rolling_replay.py b/cosmos_framework/model/generator/utils/rolling_kv/rolling_replay.py index b1e05e862..ea798a3ff 100644 --- a/cosmos_framework/model/generator/utils/rolling_kv/rolling_replay.py +++ b/cosmos_framework/model/generator/utils/rolling_kv/rolling_replay.py @@ -22,6 +22,7 @@ ) from cosmos_framework.model.generator.mot.flex_attention import CaptionMaskItem, SensorMaskItem from cosmos_framework.model.generator.utils.kv_cache import KVToStore, MultiviewARMemoryState +from cosmos_framework.model.generator.utils.memory import ARTextLayerCache from cosmos_framework.model.generator.utils.rolling_kv.rolling_prompt import ( RollingPromptLayout, RollingTextSinkLayer, @@ -49,6 +50,7 @@ class RollingReplayRequest: history: FlexQueryMetadata | None = None current: FlexQueryMetadata | None = None prompt_layout: RollingPromptLayout | None = None + text_cache_layers: list[ARTextLayerCache] | None = None def build_rolling_replay_metadata( diff --git a/cosmos_framework/model/generator/utils/rolling_kv/rolling_transfer_ar.py b/cosmos_framework/model/generator/utils/rolling_kv/rolling_transfer_ar.py index 746bdd4c7..a9f137c27 100644 --- a/cosmos_framework/model/generator/utils/rolling_kv/rolling_transfer_ar.py +++ b/cosmos_framework/model/generator/utils/rolling_kv/rolling_transfer_ar.py @@ -15,6 +15,7 @@ from cosmos_framework.utils import log from cosmos_framework.model.generator.mot.post_saturation.rolling import RollingPostSaturation from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean +from cosmos_framework.model.generator.utils.memory import ARTextLayerCache from cosmos_framework.model.generator.utils.multiview_ar import JointARChunk, joint_ar_chunks from cosmos_framework.model.generator.utils.rolling_kv.rolling_prompt import ( RollingPromptPack, @@ -189,6 +190,11 @@ def _build_branch_pack( condition_masks=condition_masks, history=cache.metadata, prompt_layout=cache.prompt_layout, + text_cache_layers=( + [ARTextLayerCache() for _ in cache.cache] + if not clean and getattr(host, "ar_use_text_cache", False) + else None + ), ) if clean: for modality in (pack.vision, pack.lidar): diff --git a/cosmos_framework/model/generator/vlm_model.py b/cosmos_framework/model/generator/vlm_model.py index 8485ce6a5..a56cbfc1d 100644 --- a/cosmos_framework/model/generator/vlm_model.py +++ b/cosmos_framework/model/generator/vlm_model.py @@ -28,16 +28,24 @@ import torch.nn as nn from torch.nn.modules.module import _IncompatibleKeys -from cosmos_framework.configs.base.defaults.parallelism import PRECISION_TO_TORCH_DTYPE -from cosmos_framework.configs.base.defaults.reasoner import validate_sound_understanding_config -from cosmos_framework.configs.base.reasoner.defaults.policy_config import VLMModelConfig +from cosmos_framework.utils.lazy_config import instantiate from cosmos_framework.model._base import ImaginaireModel +from cosmos_framework.utils import log from cosmos_framework.model.generator.algorithm.loss.cross_entropy import ( cross_entropy_loss, + weighted_ce_denominator_from_labels, weighted_cross_entropy_loss, ) from cosmos_framework.model.generator.algorithm.loss.load_balancing import compute_load_balancing_loss -from cosmos_framework.model.generator.hf_model import HFModel +from cosmos_framework.configs.base.defaults.parallelism import PRECISION_TO_TORCH_DTYPE +from cosmos_framework.configs.base.defaults.reasoner import validate_sound_understanding_config +from cosmos_framework.configs.base.reasoner.defaults.policy_config import VLMModelConfig +from cosmos_framework.model.generator.context_parallel import ( + ContextParallelBatchPolicy, + context_parallel_router_mean, + prepare_context_parallel_batch, +) +from cosmos_framework.model.generator.hf_model import HFModel, configure_sound_und_encoder_trainability from cosmos_framework.model.generator.parallelize_vlm import parallelize from cosmos_framework.model.generator.reasoner.qwen35_caption import ( Qwen35CaptionLoss, @@ -45,7 +53,7 @@ ) from cosmos_framework.model.generator.utils.moe_utils import collect_hf_moe_lbl_metadata, set_hf_moe_token_mask from cosmos_framework.model.generator.utils.safetensors_loader import load_vlm_model -from cosmos_framework.utils import log +from cosmos_framework.utils.generator.attention_chunks import plan_varlen_attention_chunks from cosmos_framework.utils.generator.input_probe import ( maybe_dump_forward_result, maybe_dump_gradients, @@ -55,14 +63,11 @@ from cosmos_framework.utils.generator.optimizer import OptimizersContainer from cosmos_framework.utils.generator.parallelism import ParallelDims from cosmos_framework.utils.generator.reasoner.constant import IGNORE_INDEX -from cosmos_framework.utils.generator.reasoner.pretrained_models_downloader import ( - maybe_download_hf_model_from_s3, -) +from cosmos_framework.utils.generator.reasoner.model_path import resolve_hf_model_path from cosmos_framework.utils.generator.reasoner.true_packing import ( TRUE_PACKING_CPU_PREPARED_KEY, assert_packing_temporal_inputs_supported, ) -from cosmos_framework.utils.lazy_config import instantiate # Model-type dispatch sets. Using hf_config.model_type (stable HF-defined string) # rather than backbone.model_name avoids the brittleness of substring-matching a local @@ -115,7 +120,15 @@ def _set_sound_und_encoder_dtype_for_fsdp( explicitly matching the encoder avoids an encoder→projector dtype mismatch. """ if hf_model.sound_und and fsdp_enabled: - hf_model.model.sound_und_model.encoder.to(dtype=PRECISION_TO_TORCH_DTYPE[precision]) + encoder = hf_model.model.sound_und_model.encoder + target_dtype = PRECISION_TO_TORCH_DTYPE[precision] + # Trainable audio parameters retain FP32 master weights under FSDP. + for parameter in encoder.parameters(): + if parameter.is_floating_point() and not parameter.requires_grad: + parameter.data = parameter.data.to(dtype=target_dtype) + for buffer in encoder.buffers(): + if buffer.is_floating_point(): + buffer.data = buffer.data.to(dtype=target_dtype) def _get_overlay_config(model_type: str) -> tuple[list[str], Callable[[str], bool]]: @@ -359,11 +372,13 @@ def __init__(self, config: VLMModelConfig, checkpoint): else: n_trainable = _apply_freeze_config(self.model.model, self.hf_config.model_type, self.config.freeze) if config.sound_und: - # The standalone artifact is the sole source of encoder weights. - # Keep it immutable even when a broad trainable_params expression - # (or a full-finetune config) would otherwise re-enable it. - self.model.model.sound_und_model.encoder.requires_grad_(False) - self.model.model.sound_und_model.encoder.eval() + # The standalone artifact supplies the initial encoder weights. + # Apply the explicit audio policy even when a broad trainable_params + # expression or a full-finetune config changes the encoder flags. + # A Stage1 checkpoint can then initialize the trainable encoder. + configure_sound_und_encoder_trainability( + self.model.model.sound_und_model, freeze_encoder=config.sound_und_config.freeze_encoder + ) if config.sound_und_config.freeze_projector: self.model.model.sound_und_model.projector.requires_grad_(False) self.model.model.sound_und_model.projector.eval() @@ -383,18 +398,9 @@ def __init__(self, config: VLMModelConfig, checkpoint): } log.info(f"freeze config applied (model_type={self.hf_config.model_type}): {self.parameter_summary}") - if self.parallel_dims is not None and self.parallel_dims.cp_enabled: - # Both CE variants normalize over every rank in the world, which is only the - # count they want while each rank holds a different sample. CP breaks that: - # its ranks hold segments of one sequence, and no reduction over those has - # been verified against the objective trained here (see the loss module - # docstring). ``_init_vlm`` already asserts cp == 1 for the attention path, - # which makes this unreachable today — it is here so the loss keeps its own - # requirement if CP is ever wired into attention. - raise NotImplementedError( - f"VLM loss does not support context parallelism (got cp={self.parallel_dims.cp}); " - "set parallelism.context_parallel_shard_degree=1." - ) + # Logits-based CE remains CP-disabled. The configured fused path uses + # CP-global sample weights and WORLD normalization, verified against CP=1. + self.context_parallel_data_window: bool = False if config.policy.use_weighted_ce: log.info(f"Using weighted CE loss with exponent={config.policy.weighted_ce_exponent}") @@ -410,16 +416,26 @@ def __init__(self, config: VLMModelConfig, checkpoint): loss_scaling_factor=1.0, ignore_index=IGNORE_INDEX, ) - # Dense Qwen3-VL and Qwen3.5 weighted CE normalize once over the whole accumulation window. - # The trainer averages microbatch losses by K; training_step therefore backprops the - # unnormalized WORLD numerator and this hook applies K / sum(global denominator). + # Opt-in Qwen3-VL and Qwen3.5 weighted CE normalize over the whole accumulation window. + # The trainer averages microbatch losses by K. CE-only windows backprop the + # unnormalized WORLD numerator and the optimizer hook applies K / sum(D). + # Active MoE auxiliary loss uses label lookahead to scale CE before backward, + # preserving the separate mean auxiliary objective without a gradient buffer. + true_pack_model_types = {"qwen3_vl", "qwen3_vl_moe"} self._window_normalize_weighted_ce = bool( config.policy.use_weighted_ce and config.policy.normalize_weighted_ce_over_accumulation_window - and self.hf_config.model_type in {"qwen3_vl", "qwen3_5"} + and self.hf_config.model_type in (true_pack_model_types | {"qwen3_5"}) ) self._weighted_ce_window_denominator: torch.Tensor | None = None self._weighted_ce_window_microbatches: int = 0 + self._weighted_ce_window_plan: torch.Tensor | None = None + self._weighted_ce_window_metrics: torch.Tensor | None = None + self.requires_accumulation_window = bool( + self._window_normalize_weighted_ce + and self.hf_config.model_type == "qwen3_vl_moe" + and config.lbl.coeff is not None + ) def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: """Initialize VLM without the legacy ModelRegistry (Phase 2+). @@ -440,7 +456,7 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: log.info(f"checkpoint.load_path: {checkpoint.load_path!r} | load_pretrain_weights: {load_pretrain_weights}") # ── a. Download HF model files (config + tokenizer; weights only if no ckpt) ── - local_path = maybe_download_hf_model_from_s3( + local_path = resolve_hf_model_path( policy.backbone.model_name, checkpoint.load_from_object_store.credentials, checkpoint.load_from_object_store.bucket, @@ -467,7 +483,12 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: configured_model_name_or_path=policy.backbone.model_name, qwen35_fp32_recurrent_a_log=policy.qwen35_fp32_recurrent_a_log, enable_fused_weighted_ce=policy.enable_fused_weighted_ce, + recompute_fused_weighted_ce=policy.recompute_fused_weighted_ce, + qwen3_vl_accumulate_head_grad_fp32=policy.qwen3_vl_accumulate_head_grad_fp32, weighted_ce_exponent=policy.weighted_ce_exponent, + qwen3_vl_moe_fixed_sum=policy.qwen3_vl_moe_fixed_sum, + qwen3_vl_vision_attention_max_padded_tokens=policy.qwen3_vl_vision_attention_max_padded_tokens, + qwen3_vl_decoder_attention_max_padded_tokens=policy.qwen3_vl_decoder_attention_max_padded_tokens, ) from cosmos_framework.model.generator.qwen3_vl_compat import apply_qwen3_vl_patch_embed_compat @@ -507,7 +528,7 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: # ── c. Build ParallelDims + device mesh ── # Overlay-mesh design (see vfm/utils/parallelism.py): cp/cfgp do NOT # consume FSDP rank slots, so dp_replicate * dp_shard == world_size - # alone. The VLM HFModel doesn't have a CP-aware attention path. + # alone. Qwen3-VL uses a Ulysses CP overlay on that same world. world_size = int(os.environ.get("WORLD_SIZE", 1)) _dp_replicate = config.parallelism.data_parallel_replicate_degree # Single-process run: force dp_replicate=1 so ParallelDims doesn't @@ -524,12 +545,17 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: enable_inference_mode=False, ) - # VLM does not currently support cp or cfgp. CP needs a CP-aware - # attention path (see ``vfm/models/mot/context_parallel_utils.py``) that - # is not wired into the VLM HFModel; CFGP is inference-only. - assert parallel_dims.cp == 1, f"VLM does not support CP (got cp={parallel_dims.cp})" + # CP is opt-in for validated native fused-loss integrations; Qwen multimodal encoding has a separate flag. + # The CP-aware attention path is wired before AC/FSDP wrap the blocks; + # CFGP remains inference-only and unsupported here. assert parallel_dims.cfgp == 1, f"VLM does not support CFGP (got cfgp={parallel_dims.cfgp})" - + if parallel_dims.cp_enabled: + if config.sound_und or config.compile.enabled: + raise NotImplementedError("Qwen3-VL CP requires sound_und=False and compile=False") + if config.lbl.coeff is not None and config.lbl.method != "global": + raise NotImplementedError("Initial Qwen3-VL CP requires global load balancing") + if not policy.enable_fused_weighted_ce: + raise ValueError("Qwen3-VL CP requires enable_fused_weighted_ce=True") if torch.distributed.is_initialized(): parallel_dims.build_meshes(device_type="cuda") @@ -544,11 +570,11 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: fsdp_enabled=parallel_dims.dp_enabled, ) - # ── d. Apply activation checkpointing, torch.compile and FSDP2 ── + # ── d. Apply CP, activation checkpointing, torch.compile and FSDP2 ── # config.compile is threaded through so model.config.compile.enabled=True # actually compiles each block in place (was previously a dead config on # the VLM path — only the MoT path consumed it). See parallelize_vlm for why the - # three passes must run in this order (AC wraps, compile compiles the wrapper, FSDP + # wrapping passes must run after CP setup in this order (AC wraps, compile compiles the wrapper, FSDP # shards it). # Called unconditionally, NOT under an is_initialized() guard: activation # checkpointing and compile are independent of distribution, and a single-process @@ -562,6 +588,7 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: config.precision, activation_checkpointing=config.activation_checkpointing, compile_config=config.compile, + allow_multimodal_cp=policy.qwen3_vl_cp_multimodal, ) # ── e. Materialize meta tensors on CUDA ── @@ -588,7 +615,7 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: # ── g. Load pretrain weights ── if load_pretrain_weights: if policy.backbone.safetensors_path: - safetensors_local_path = maybe_download_hf_model_from_s3( + safetensors_local_path = resolve_hf_model_path( policy.backbone.safetensors_path, checkpoint.load_from_object_store.credentials, checkpoint.load_from_object_store.bucket, @@ -617,7 +644,7 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: if llm_path: overlay_skip_patterns, is_lm_key = _get_overlay_config(hf_model.hf_config.model_type) - llm_local_path = maybe_download_hf_model_from_s3( + llm_local_path = resolve_hf_model_path( llm_path, checkpoint.load_from_object_store.credentials, checkpoint.load_from_object_store.bucket, @@ -682,6 +709,38 @@ def on_after_backward(self, iteration: int = 0) -> None: """Capture exact pre-clip gradients when deep parity probing is enabled.""" maybe_dump_gradients(self.model.model, self._parity_probe_step, tag="i4") + def prepare_accumulation_window(self, batches: tuple[dict[str, Any], ...]) -> None: + """Know the full CE denominator before mixing CE and auxiliary gradients. + + With G microbatches the objective is sum(N_g)/sum(D_g) + mean(A_g). + CPU lookahead avoids a second full-model gradient buffer or retained graphs. + Only the scalar denominators move to the device; batches remain on CPU. + """ + if not self.requires_accumulation_window or len(batches) < 2: + raise ValueError("MoE accumulation lookahead requires an enabled objective and GAS > 1") + if self._weighted_ce_window_plan is not None or self._weighted_ce_window_denominator is not None: + raise RuntimeError("The preceding weighted-CE accumulation window has not finished") + denominators: list[torch.Tensor] = [] + for batch in batches: + labels = batch.get("labels") + if not isinstance(labels, torch.Tensor): + raise TypeError("MoE accumulation lookahead requires pretokenized CPU labels") + boundaries = batch.get("packed_cu_seq_lens") + if batch.get("true_packing") and not isinstance(boundaries, torch.Tensor): + raise ValueError("Packed accumulation lookahead requires conversation boundaries") + denominators.append( + weighted_ce_denominator_from_labels( + labels, exponent=self.config.policy.weighted_ce_exponent, cu_seq_lens=boundaries + ) + ) + plan = torch.stack(denominators).to(next(self.parameters()).device) + if self.parallel_dims is not None and self.parallel_dims.cp_enabled: + # The CP-aware loader replicates full labels, unlike fused CE's local targets. + plan.div_(self.parallel_dims.cp_mesh.size()) + if torch.distributed.is_initialized(): + torch.distributed.all_reduce(plan) + self._weighted_ce_window_plan = plan + def on_before_optimizer_step( self, optimizer: torch.optim.Optimizer | OptimizersContainer, @@ -694,15 +753,23 @@ def on_before_optimizer_step( return if self._weighted_ce_window_microbatches <= 0: raise RuntimeError("weighted-CE denominator exists without accumulated microbatches") - scale = self._weighted_ce_window_microbatches / self._weighted_ce_window_denominator.clamp(min=1) - optimizers = optimizer.optimizers if isinstance(optimizer, OptimizersContainer) else [optimizer] - for inner_optimizer in optimizers: - for group in inner_optimizer.param_groups: - for parameter in group["params"]: - if parameter.grad is not None: - parameter.grad.mul_(scale) + if self._weighted_ce_window_plan is not None: + if self._weighted_ce_window_microbatches != self._weighted_ce_window_plan.numel(): + raise RuntimeError("Optimizer step before the complete planned accumulation window") + # CE was scaled before backward; rescaling the combined gradients here + # would incorrectly change the independently averaged auxiliary objective. + self._weighted_ce_window_plan = None + else: + scale = self._weighted_ce_window_microbatches / self._weighted_ce_window_denominator.clamp(min=1) + optimizers = optimizer.optimizers if isinstance(optimizer, OptimizersContainer) else [optimizer] + for inner_optimizer in optimizers: + for group in inner_optimizer.param_groups: + for parameter in group["params"]: + if parameter.grad is not None: + parameter.grad.mul_(scale) self._weighted_ce_window_denominator = None self._weighted_ce_window_microbatches = 0 + self._weighted_ce_window_metrics = None def on_before_zero_grad( self, @@ -826,15 +893,24 @@ def _moe_load_balancing_loss(self) -> torch.Tensor | None: if self.hf_config.model_type != "qwen3_vl_moe": return None + metadata = collect_hf_moe_lbl_metadata(self.model) + if metadata is not None and self.parallel_dims.cp_enabled: + # Recover each full sequence's router mean when CP shards have unequal + # valid lengths (right padding). WORLD averaging then matches CP=1. + metadata = metadata._replace( + mean_router_prob_per_expert=context_parallel_router_mean( + metadata.mean_router_prob_per_expert, metadata.num_tokens, self.parallel_dims.cp_mesh.get_group() + ) + ) return compute_load_balancing_loss( - collect_hf_moe_lbl_metadata(self.model), + metadata, coeff=self.config.lbl.coeff, method=self.config.lbl.method, device_mesh=self.parallel_dims.dp_mesh if self.parallel_dims is not None else None, ) def _prepare_true_packing(self, data: dict[str, Any]) -> None: - """Prepare dense Qwen3-VL true-packed positions and standard varlen metadata in place. + """Prepare shared true-packed positions and model-family varlen metadata in place. Padded batches return immediately and continue using the current model-internal M-RoPE path. Packed batches are accepted only for the one backend/model combination whose @@ -845,10 +921,18 @@ def _prepare_true_packing(self, data: dict[str, Any]) -> None: return if true_packing is not True: raise TypeError("true_packing must be a Python bool") - if self.hf_config.model_type != "qwen3_vl": + cp_size = self.parallel_dims.cp_size if self.parallel_dims is not None else 1 + loader_cp_size = data.get("context_parallel_size", 1) + if loader_cp_size != cp_size: + raise ValueError(f"Loader CP={loader_cp_size} must match model CP={cp_size}") + supported_model_types = {"qwen3_vl", "qwen3_vl_moe"} + if self.hf_config.model_type not in supported_model_types: raise NotImplementedError( - f"True packing is validated only for dense qwen3_vl; got model_type={self.hf_config.model_type!r}" + f"True packing is validated only for {sorted(supported_model_types)}; " + f"got model_type={self.hf_config.model_type!r}" ) + if self.hf_config.model_type == "qwen3_vl_moe" and not self.config.policy.enable_fused_weighted_ce: + raise ValueError("Qwen3-VL MoE packing requires fused weighted CE") if self.config.sound_und: raise NotImplementedError( "True packing is not validated for audio understanding inputs; use padded batching." @@ -898,11 +982,28 @@ def _prepare_true_packing(self, data: dict[str, Any]) -> None: ) # Assign the same cumulative-boundary tensor object to Q and K after its one H2D copy. - # The cosmos adapter accepts the standard HF protocol directly. + # Qwen attention and weighted CE consume these same conversation boundaries. data["cu_seq_lens_q"] = packed_cu_seq_lens data["cu_seq_lens_k"] = packed_cu_seq_lens data["max_length_q"] = packed_max_length data["max_length_k"] = packed_max_length + + budget = getattr(self.config.policy, "qwen3_vl_decoder_attention_max_padded_tokens", 0) + if budget: + if self.hf_config.model_type not in {"qwen3_vl", "qwen3_vl_moe"}: + raise ValueError("Bounded decoder varlen attention requires Qwen3-VL") + lengths = list(seq_lens) + padding = input_ids.shape[1] - sum(lengths) + if padding < 0: + raise ValueError("Packed logical lengths exceed the physical row") + if padding: + lengths.append(padding) + if packed_cu_seq_lens.numel() != len(lengths) + 1: + raise ValueError("Packed boundaries must include every conversation and the padding tail") + # Plan once from CPU integers, before CP shards the decoder sequence. + # Each attention head owner applies these global boundaries after all-to-all. + data["decoder_attention_chunks"] = plan_varlen_attention_chunks(lengths, budget) + data.pop("true_packing", None) def training_step(self, data: dict, iteration: int) -> tuple[dict, torch.Tensor]: @@ -937,8 +1038,12 @@ def training_step(self, data: dict, iteration: int) -> tuple[dict, torch.Tensor] ``unit_tests/test_monkey_patch.py`` pins it. """ self._prepare_true_packing(data) + if self.parallel_dims is not None and self.parallel_dims.cp_enabled: + cp_policy = self.model.context_parallel_batch_policy + assert isinstance(cp_policy, ContextParallelBatchPolicy) + data = prepare_context_parallel_batch(data, self.parallel_dims.cp_mesh.get_group(), policy=cp_policy) labels = data.pop("labels") - self._set_moe_token_mask(data.get("attention_mask")) + self._set_moe_token_mask(data.get("context_parallel_token_mask", data.get("attention_mask"))) maybe_dump_model_inputs(data, iteration, tag="i4", labels=labels) self._parity_probe_step = iteration @@ -966,6 +1071,13 @@ def training_step(self, data: dict, iteration: int) -> tuple[dict, torch.Tensor] if self._window_normalize_weighted_ce: denominator = loss_stats.global_objective_denominator backward_loss = loss * denominator + plan = self._weighted_ce_window_plan + if plan is not None: + index = self._weighted_ce_window_microbatches + if index >= plan.numel() or not torch.isclose(denominator, plan[index], rtol=1e-5, atol=1e-5): + raise RuntimeError("Forward CE denominator differs from the CPU accumulation plan") + # The trainer divides each backward loss by G. + backward_loss = backward_loss * (plan.numel() / plan.sum().clamp(min=1)) if self._weighted_ce_window_denominator is None: self._weighted_ce_window_denominator = denominator.detach().clone() else: @@ -978,7 +1090,32 @@ def training_step(self, data: dict, iteration: int) -> tuple[dict, torch.Tensor] load_balancing_loss = self._moe_load_balancing_loss() if load_balancing_loss is not None: loss = loss + load_balancing_loss - backward_loss = backward_loss + load_balancing_loss + if self._weighted_ce_window_plan is not None: + backward_loss = backward_loss + load_balancing_loss + elif self._window_normalize_weighted_ce: + if self._weighted_ce_window_microbatches != 1: + raise NotImplementedError( + "Without prepare_accumulation_window, window-normalized MoE CE requires " + "grad_accum_iter=1. Use the trainer's CPU lookahead for GAS > 1; rescaling " + "combined gradients would change the auxiliary objective." + ) + # The optimizer hook divides by this same denominator. Keep the + # auxiliary coefficient unchanged for the supported single-step window. + backward_loss = backward_loss + load_balancing_loss * loss_stats.global_objective_denominator.clamp( + min=1 + ) + else: + backward_loss = backward_loss + load_balancing_loss + if self._weighted_ce_window_plan is not None: + plan = self._weighted_ce_window_plan + auxiliary = torch.zeros_like(ce_loss) if load_balancing_loss is None else load_balancing_loss.detach() + metrics = torch.stack( + (ce_loss * loss_stats.global_objective_denominator / plan.sum().clamp(min=1), auxiliary / plan.numel()) + ) + if self._weighted_ce_window_metrics is None: + self._weighted_ce_window_metrics = metrics + else: + self._weighted_ce_window_metrics.add_(metrics) if not isinstance(logits, Qwen35CaptionLoss): # The fused path deliberately never creates vocabulary logits for a probe. maybe_dump_forward_result(logits, {"ce_loss": ce_loss, "total_loss": loss}, iteration, tag="i4") @@ -1009,6 +1146,15 @@ def training_step(self, data: dict, iteration: int) -> tuple[dict, torch.Tensor] "per_sample_objective_weight": loss_stats.per_sample_objective_denominator, "global_objective_weight": loss_stats.global_objective_denominator, } + if ( + self._weighted_ce_window_plan is not None + and self._weighted_ce_window_microbatches == self._weighted_ce_window_plan.numel() + ): + # The official logger observes scalar "*loss" keys at optimizer boundaries. + # Publish the complete window without changing legacy microbatch loss keys. + assert self._weighted_ce_window_metrics is not None + output["window_ce_loss"], output["window_aux_loss"] = self._weighted_ce_window_metrics.unbind() + output["window_total_loss"] = self._weighted_ce_window_metrics.sum() if backward_loss is not loss: output["_backward_loss"] = backward_loss if load_balancing_loss is not None: @@ -1027,8 +1173,12 @@ def validation_step(self, data: dict, iteration: int) -> tuple[dict, torch.Tenso ``attention_mask`` is forwarded (see that method's notes). """ self._prepare_true_packing(data) + if self.parallel_dims is not None and self.parallel_dims.cp_enabled: + cp_policy = self.model.context_parallel_batch_policy + assert isinstance(cp_policy, ContextParallelBatchPolicy) + data = prepare_context_parallel_batch(data, self.parallel_dims.cp_mesh.get_group(), policy=cp_policy) labels = data.pop("labels") - self._set_moe_token_mask(data.get("attention_mask")) + self._set_moe_token_mask(data.get("context_parallel_token_mask", data.get("attention_mask"))) forward_loss_kwargs = ( {"labels": labels} if getattr(self.config.policy, "enable_fused_weighted_ce", False) else {} ) diff --git a/cosmos_framework/scripts/train_experiment.py b/cosmos_framework/scripts/train_experiment.py index 5350c6597..b30f70109 100644 --- a/cosmos_framework/scripts/train_experiment.py +++ b/cosmos_framework/scripts/train_experiment.py @@ -114,11 +114,21 @@ def _force_compile_disabled(cfg) -> int: def launch(config: Config, args: argparse.Namespace) -> None: + if os.environ.get("COSMOS_CLUSTERPULSE_ENABLED") == "1" and ( + config.trainer.profiling.enable_nsys or config.trainer.profiling.enable_profiling + ): + raise ValueError("ClusterPulse cannot be combined with nsys/PyTorch profiling in the resolved experiment") + from cosmos_framework.utils.lazy_config import instantiate from cosmos_framework.utils import distributed from cosmos_framework.utils.context_managers import data_loader_init, distributed_init, model_init from cosmos_framework.utils.launch import log_reproducible_setup + if os.environ.get("COSMOS_CLUSTERPULSE_HEALTH_CHECK"): + from cosmos_framework.utils.generator.clusterpulse_preflight import prepare_clusterpulse_providers + + prepare_clusterpulse_providers() + # The released configs name their tokenizer and VAE by object-store URI; the framework's # checkpoint registry maps those URIs to public downloads. from cosmos_framework.inference.common.checkpoints import register_checkpoints diff --git a/cosmos_framework/trainer/__init__.py b/cosmos_framework/trainer/__init__.py index 3cf90a684..797629904 100644 --- a/cosmos_framework/trainer/__init__.py +++ b/cosmos_framework/trainer/__init__.py @@ -5,7 +5,8 @@ import inspect import os import signal -from collections.abc import Iterator +from collections import deque +from collections.abc import Callable, Iterator from dataclasses import dataclass, field from typing import Any @@ -13,8 +14,8 @@ import torch.distributed as dist import torch.utils.data -from cosmos_framework.utils.context_managers import distributed_init from cosmos_framework.utils.flags import INTERNAL +from cosmos_framework.utils.context_managers import distributed_init from cosmos_framework.utils.profiling import ( maybe_enable_memory_snapshot, maybe_enable_nsys_profiling, @@ -29,13 +30,14 @@ USE_MEGATRON = False +from cosmos_framework.utils.lazy_config import LazyConfig, instantiate from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.utils import callback, distributed, ema, log, misc from cosmos_framework.utils.checkpointer import Checkpointer -from cosmos_framework.utils.lazy_config import LazyConfig, instantiate from cosmos_framework.utils.misc import StragglerDetectorV2 + @dataclass class ContextParallelDataWindow: """Caches one dataloader batch across a ``cp_size``-step CP data window. @@ -141,6 +143,42 @@ def assert_synced_with_model(self, model: Any) -> None: ) +@dataclass +class AccumulationDataWindow: + """Bounded CPU lookahead; consumed batches are released before the next fetch. + + An incomplete window survives an epoch boundary. No optimizer checkpoint can + observe unread lookahead, because all batches are consumed before the update. + """ + + batches: deque[dict[str, Any]] = field(default_factory=deque) + ready: bool = False + last_stop_count: int | None = None + + def fetch( + self, + iterator: Iterator[dict[str, Any]], + size: int, + prepare: Callable[[tuple[dict[str, Any], ...]], None], + ) -> tuple[dict[str, Any] | None, bool]: + if not self.ready: + try: + while len(self.batches) < size: + self.batches.append(next(iterator)) + except StopIteration: + if self.last_stop_count == len(self.batches): + raise RuntimeError("Cannot complete an accumulation window from an empty dataloader") + self.last_stop_count = len(self.batches) + return None, True + prepare(tuple(self.batches)) + self.ready = True + self.last_stop_count = None + batch = self.batches.popleft() + if not self.batches: + self.ready = False + return batch, False + + @dataclass class StepBatch: """Encapsulates batch routing for one training step. @@ -350,6 +388,7 @@ def __init__(self, config): self.straggler_detector.initialize() # One fetched batch reused for cp_size consecutive steps when CP is enabled. self._cp_data_window = ContextParallelDataWindow() + self._accumulation_data_window = AccumulationDataWindow() # Send a TimeoutError if a training step takes over timeout_period seconds. signal.signal(signal.SIGALRM, functools.partial(misc.timeout_handler, config.trainer.timeout_period)) # type: ignore @@ -361,7 +400,8 @@ def _fetch_data_batch( """ Fetch the next training batch. - Without CP, every rank fetches a new batch for each step. With CP, each rank + Without CP, or with a loader that already replicates CP samples, fetch a + new batch each step. Other CP models retain their rotation window: each rank fetches one independent local batch at the start of a ``cp_size``-step data window. ``self._cp_data_window`` returns that same batch for every slot (each slot is one training step within the window; CP rank ``s`` owns slot @@ -373,15 +413,26 @@ def _fetch_data_batch( - data_batch: The fetched data batch (or None if stopped). - stop_signal (bool): True if StopIteration was encountered. """ + if getattr(model, "requires_accumulation_window", False) and self.config.trainer.grad_accum_iter > 1: + if self._context_parallel_size(model) != 1: + raise ValueError("Accumulation lookahead requires a loader that already replicates CP samples") + return self._accumulation_data_window.fetch( + dataloader_iter, self.config.trainer.grad_accum_iter, model.prepare_accumulation_window + ) parallel_dims = getattr(model, "parallel_dims", None) - if parallel_dims is None or not parallel_dims.cp_enabled: + if ( + parallel_dims is None + or not parallel_dims.cp_enabled + or not getattr(model, "context_parallel_data_window", True) + ): try: return next(dataloader_iter), False except StopIteration: return None, True cp_size = parallel_dims.cp_mesh.size() - self._cp_data_window.assert_synced_with_model(model) + if not getattr(self.config.trainer, "skip_model_update", False): + self._cp_data_window.assert_synced_with_model(model) if not self._cp_data_window.active: try: self._cp_data_window.batch = next(dataloader_iter) @@ -489,7 +540,11 @@ def _release_cp_raw_batch_after_preprocessing(model: ImaginaireModel) -> bool: @staticmethod def _context_parallel_size(model: ImaginaireModel) -> int: parallel_dims = getattr(model, "parallel_dims", None) - if parallel_dims is None or not parallel_dims.cp_enabled: + if ( + parallel_dims is None + or not parallel_dims.cp_enabled + or not getattr(model, "context_parallel_data_window", True) + ): return 1 return int(parallel_dims.cp_mesh.size()) @@ -501,6 +556,7 @@ def _resume_dataloader_fetch_count(self, model: ImaginaireModel, iteration: int) the next raw batch avoids replaying any data from that partial window. """ microsteps = iteration * self.config.trainer.grad_accum_iter + # Models with an already CP-replicated loader consume one batch each microstep. cp_size = self._context_parallel_size(model) return (microsteps + cp_size - 1) // cp_size @@ -521,10 +577,18 @@ def train( model = model.to("cuda", memory_format=self.config.trainer.memory_format) # type: ignore model.on_train_start(self.config.trainer.memory_format) + validate_training_horizon = getattr(dataloader_train, "validate_training_horizon", None) + if callable(validate_training_horizon): + validate_training_horizon(self._resume_dataloader_fetch_count(model, self.max_iterations)) + # Initialize the optimizer, scheduler, and grad_scaler. self.callbacks.on_optimizer_init_start() optimizer, scheduler = model.init_optimizer_scheduler(self.config.optimizer, self.config.scheduler) grad_scaler = torch.amp.GradScaler("cuda", **self.config.trainer.grad_scaler_args) + # Checkpoint callbacks also need these when optimizer-step hooks are skipped. + self.optimizer: torch.optim.Optimizer = optimizer + self.scheduler: torch.optim.lr_scheduler.LRScheduler = scheduler + self.grad_scaler: torch.amp.GradScaler = grad_scaler self.callbacks.on_optimizer_init_end() # Load the model checkpoint and get the starting iteration number. iteration = self.checkpointer.load(model, optimizer, scheduler, grad_scaler) @@ -604,15 +668,20 @@ def train( model_ddp.train() assert model_ddp.training, "model_ddp is not in training mode." assert model.training, "model is not in training mode." - output_batch, loss, grad_accum_iter = self.training_step( - model_ddp, - optimizer, - scheduler, - grad_scaler, - step_batch.training, - iteration=iteration, - grad_accum_iter=grad_accum_iter, - ) + if self.config.trainer.skip_model_update: + output_batch = {} + loss = torch.zeros((), device="cuda") # [] + grad_accum_iter = (grad_accum_iter + 1) % self.config.trainer.grad_accum_iter + else: + output_batch, loss, grad_accum_iter = self.training_step( + model_ddp, + optimizer, + scheduler, + grad_scaler, + step_batch.training, + iteration=iteration, + grad_accum_iter=grad_accum_iter, + ) step_batch.finalize() self.callbacks.on_training_step_batch_end( model, step_batch.callback_post_step, output_batch, loss, iteration=iteration @@ -622,7 +691,7 @@ def train( # Release this microstep's outputs before the next forward/backward. del output_batch, loss continue - # Do the following when an actual optimizer (update) step has been made. + # Finish an optimizer step, or its equivalent dataloading window in skip mode. iteration += 1 # Save checkpoint. completed_epoch = ( @@ -638,13 +707,10 @@ def train( ) iter_save_due = save_freq_in_epoch == 0 and iteration % self.config.checkpoint.save_iter == 0 if epoch_save_due or iter_save_due: + # Only epoch saves pass ``epoch``; not every checkpointer accepts it. + save_kwargs = {"epoch": completed_epoch} if epoch_save_due else {} self.checkpointer.save( - model, - optimizer, - scheduler, - grad_scaler, - iteration=iteration, - epoch=completed_epoch if epoch_save_due else None, + model, optimizer, scheduler, grad_scaler, iteration=iteration, **save_kwargs ) self.callbacks.on_training_step_end( model, @@ -663,10 +729,10 @@ def train( and validation_freq_in_epoch > 0 and completed_epoch % validation_freq_in_epoch == 0 ) - iter_validation_due = ( - validation_freq_in_epoch == 0 and iteration % self.config.trainer.validation_iter == 0 - ) - if self.config.trainer.run_validation and (epoch_validation_due or iter_validation_due): + if self.config.trainer.run_validation and ( + epoch_validation_due + or (validation_freq_in_epoch == 0 and iteration % self.config.trainer.validation_iter == 0) + ): self.validate(model, dataloader_val, iteration=iteration) # This iteration is successful; reset the timeout signal. signal.alarm(self.config.trainer.timeout_period) @@ -693,14 +759,8 @@ def train( final_epoch is not None and save_freq_in_epoch > 0 and final_epoch % save_freq_in_epoch == 0 ) or (save_freq_in_epoch == 0 and iteration % self.config.checkpoint.save_iter == 0) if not final_already_saved: - self.checkpointer.save( - model, - optimizer, - scheduler, - grad_scaler, - iteration=iteration, - epoch=final_epoch if save_freq_in_epoch > 0 else None, - ) + save_kwargs = {"epoch": final_epoch} if save_freq_in_epoch > 0 else {} + self.checkpointer.save(model, optimizer, scheduler, grad_scaler, iteration=iteration, **save_kwargs) self.callbacks.on_train_end(model, iteration=iteration) self.checkpointer.finalize() distributed.barrier() diff --git a/cosmos_framework/trainer/distillation.py b/cosmos_framework/trainer/distillation.py index 588c5a71e..4fa60a2f3 100644 --- a/cosmos_framework/trainer/distillation.py +++ b/cosmos_framework/trainer/distillation.py @@ -9,13 +9,13 @@ import torch from loguru import logger as log +from cosmos_framework.trainer import ImaginaireTrainer +from cosmos_framework.utils import distributed from cosmos_framework.model.generator.distillation.optimizer import ( PhaseOptimizer, PhaseScheduler, iter_torch_optimizers, ) -from cosmos_framework.trainer import ImaginaireTrainer -from cosmos_framework.utils import distributed __all__: tuple[str, ...] = ("DistillationTrainer",) diff --git a/cosmos_framework/trainer/distillation_test.py b/cosmos_framework/trainer/distillation_test.py index a9fe15aee..68377dad4 100644 --- a/cosmos_framework/trainer/distillation_test.py +++ b/cosmos_framework/trainer/distillation_test.py @@ -10,9 +10,9 @@ import torch import cosmos_framework.trainer.distillation as trainer_module -from cosmos_framework.model.generator.distillation.optimizer import PhaseOptimizer, PhaseScheduler from cosmos_framework.trainer import ImaginaireTrainer from cosmos_framework.trainer.distillation import DistillationTrainer +from cosmos_framework.model.generator.distillation.optimizer import PhaseOptimizer, PhaseScheduler from cosmos_framework.utils.generator.optimizer import OptimizersContainer diff --git a/cosmos_framework/utils/callback.py b/cosmos_framework/utils/callback.py index 8a922ed0b..fb1b4a746 100644 --- a/cosmos_framework/utils/callback.py +++ b/cosmos_framework/utils/callback.py @@ -15,8 +15,8 @@ import tqdm import wandb -from cosmos_framework.utils import distributed, log, misc, wandb_util from cosmos_framework.utils.lazy_config import LazyCall, instantiate +from cosmos_framework.utils import distributed, log, misc, wandb_util from cosmos_framework.utils.misc import get_local_tensor_if_DTensor try: @@ -26,9 +26,9 @@ if TYPE_CHECKING: + from cosmos_framework.utils.config import Config from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.trainer import ImaginaireTrainer - from cosmos_framework.utils.config import Config class CallBackGroup: @@ -94,6 +94,8 @@ def __getattr__(self, method_name: str) -> Callable: def multi_callback_wrapper(*args, **kwargs) -> None: for callback in self._callbacks: + if method_name in _OPTIONAL_CALLBACK_HOOKS and not callable(getattr(callback, method_name, None)): + continue assert hasattr(callback, method_name) method = getattr(callback, method_name) assert callable(method) @@ -316,12 +318,20 @@ def on_exception(self, error: BaseException) -> None: pass +# Hooks added after the callback interface was established. Callbacks that do not subclass +# Callback may omit them; CallBackGroup then skips them for that callback. +_OPTIONAL_CALLBACK_HOOKS = frozenset({"on_exception"}) + + def _missing_callback_hooks(callback: object) -> list[str]: """Return required callback hooks that are absent or non-callable.""" return sorted( name for name, member in vars(Callback).items() - if name.startswith("on_") and callable(member) and not callable(getattr(callback, name, None)) + if name.startswith("on_") + and name not in _OPTIONAL_CALLBACK_HOOKS + and callable(member) + and not callable(getattr(callback, name, None)) ) diff --git a/cosmos_framework/utils/config.py b/cosmos_framework/utils/config.py index ef7fe44f7..b9a6a710e 100644 --- a/cosmos_framework/utils/config.py +++ b/cosmos_framework/utils/config.py @@ -29,9 +29,9 @@ USE_MEGATRON = False -from cosmos_framework.utils import distributed from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.utils.lazy_config import LazyDict +from cosmos_framework.utils import distributed from cosmos_framework.utils.misc import Color T = TypeVar("T") @@ -487,6 +487,8 @@ class TrainerConfig: memory_format: torch.memory_format = torch.preserve_format # Gradient accumulation (update step every N iteration). grad_accum_iter: int = 1 + # Debug dataloading without training forward/backward or optimizer/scheduler updates. + skip_model_update: bool = False # Straggler Detection config straggler_detection: StragglerDetectionConfig = attrs.field(factory=StragglerDetectionConfig) # Profiling config diff --git a/cosmos_framework/utils/distributed.py b/cosmos_framework/utils/distributed.py index ccf8f0ad5..569f37a2d 100644 --- a/cosmos_framework/utils/distributed.py +++ b/cosmos_framework/utils/distributed.py @@ -23,8 +23,8 @@ import torch.distributed as dist from torch.distributed import get_process_group_ranks -from cosmos_framework.utils.device import Device from cosmos_framework.utils.flags import INTERNAL +from cosmos_framework.utils.device import Device if dist.is_available(): from torch.distributed.distributed_c10d import _get_default_group diff --git a/cosmos_framework/utils/easy_io/backends/__init__.py b/cosmos_framework/utils/easy_io/backends/__init__.py index 5211dedb3..99c8886cd 100644 --- a/cosmos_framework/utils/easy_io/backends/__init__.py +++ b/cosmos_framework/utils/easy_io/backends/__init__.py @@ -19,13 +19,17 @@ if TRAINING: try: from cosmos_framework.utils.easy_io.backends.boto3_backend import Boto3Backend - except ImportError: - pass + except ModuleNotFoundError as error: + # Only a missing optional S3 client is tolerated; any other import failure surfaces here. + if (error.name or "").split(".")[0] not in {"boto3", "botocore", "multistorageclient"}: + raise else: __all__.append("Boto3Backend") try: from cosmos_framework.utils.easy_io.backends.msc_backend import MSCBackend - except ImportError: - pass + except ModuleNotFoundError as error: + # Only a missing optional S3 client is tolerated; any other import failure surfaces here. + if (error.name or "").split(".")[0] not in {"boto3", "botocore", "multistorageclient"}: + raise else: __all__.append("MSCBackend") diff --git a/cosmos_framework/utils/easy_io/backends/registry_utils.py b/cosmos_framework/utils/easy_io/backends/registry_utils.py index 11581cf55..e446892f7 100644 --- a/cosmos_framework/utils/easy_io/backends/registry_utils.py +++ b/cosmos_framework/utils/easy_io/backends/registry_utils.py @@ -117,10 +117,13 @@ def _register(backend_cls): if TRAINING: try: from cosmos_framework.utils.easy_io.backends.msc_backend import MSCBackend - except ImportError as error: + except ModuleNotFoundError as error: # multistorageclient is an optional internal S3 dependency. Local, # HTTP, Docker, and Pyxis training must remain importable without it. # A later s3:// access fails clearly because no S3 backend is present. + # Any other import failure (a bug in the backend or its imports) surfaces here. + if (error.name or "").split(".")[0] not in {"boto3", "botocore", "multistorageclient"}: + raise MSC_IMPORT_ERROR = error else: # To avoid breaking backward Compatibility, 's3' is also used as a diff --git a/cosmos_framework/utils/easy_io/file_client.py b/cosmos_framework/utils/easy_io/file_client.py index cb3af4c54..9c9acd4af 100644 --- a/cosmos_framework/utils/easy_io/file_client.py +++ b/cosmos_framework/utils/easy_io/file_client.py @@ -75,8 +75,10 @@ class FileClient: if TRAINING: try: from cosmos_framework.utils.easy_io.backends.msc_backend import MSCBackend - except ImportError: - pass + except ModuleNotFoundError as error: + # Only a missing optional S3 client is tolerated; any other import failure surfaces here. + if (error.name or "").split(".")[0] not in {"boto3", "botocore", "multistorageclient"}: + raise else: _backends["s3"] = MSCBackend _backends["msc"] = MSCBackend diff --git a/cosmos_framework/utils/generator/attention_chunks.py b/cosmos_framework/utils/generator/attention_chunks.py new file mode 100644 index 000000000..1fe44bf67 --- /dev/null +++ b/cosmos_framework/utils/generator/attention_chunks.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""CPU plans for bounded varlen workspaces without splitting independent sequences.""" + + +def plan_varlen_attention_chunks( + lengths: list[int], max_padded_tokens: int +) -> tuple[tuple[int, int, int, int, int], ...]: + """Bound sequence count times maximum length per attention call. + + Entries are (first_sequence, end_sequence, first_token, end_token, max_length). + Preserve sequence order and boundaries, including any isolated padding tail. + A single sequence longer than the budget stays intact. + """ + if max_padded_tokens <= 0 or not lengths or any(length <= 0 for length in lengths): + raise ValueError("Sequence lengths and the padded-token budget must be positive") + chunks: list[tuple[int, int, int, int, int]] = [] + first_sequence = first_token = end_token = max_length = 0 + for sequence, length in enumerate(lengths): + proposed_max = max(max_length, length) + if sequence > first_sequence and (sequence - first_sequence + 1) * proposed_max > max_padded_tokens: + chunks.append((first_sequence, sequence, first_token, end_token, max_length)) + first_sequence, first_token, max_length = sequence, end_token, 0 + end_token += length + max_length = max(max_length, length) + chunks.append((first_sequence, len(lengths), first_token, end_token, max_length)) + return tuple(chunks) diff --git a/cosmos_framework/utils/generator/clusterpulse_preflight.py b/cosmos_framework/utils/generator/clusterpulse_preflight.py new file mode 100644 index 000000000..bba85193c --- /dev/null +++ b/cosmos_framework/utils/generator/clusterpulse_preflight.py @@ -0,0 +1,67 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Check ClusterPulse GPU and in-process NVTX providers before distributed training.""" + +import fcntl +import os +import socket +import subprocess +import time +from pathlib import Path + + +def emit_nvtx_probe() -> None: + # Import only on the opt-in training path; ordinary launches keep their import order. + import torch + + torch.cuda.set_device(int(os.environ["SLURM_LOCALID"])) + torch.cuda.init() + with torch.cuda.nvtx.range(f"clusterpulse_provider_probe_rank_{os.environ['SLURM_PROCID']}"): + torch.cuda.synchronize() + + +def run_health_check(checker: Path, *, process: bool) -> None: + env = dict(os.environ) + # The observer uses the same compat UMD but does not need workload injection. + env.pop("LD_PRELOAD", None) + env["NVAON_LOG_ENABLE_STDERR"] = "1" + command = [str(checker), "--service-port", os.environ["NVAON_SERVICE_PORT"], "--skip-pm"] + timeout = 10 if process else 60 + command.extend(["--timeout", str(timeout)]) + if process: + command.extend(["--process-pid", str(os.getpid())]) + try: + subprocess.run(command, env=env, check=True, timeout=timeout + 5) + except (OSError, subprocess.SubprocessError) as error: + scope = "workload NVTX" if process else "GPU trace" + raise RuntimeError(f"ClusterPulse {scope} provider readiness failed before training") from error + + +def prepare_clusterpulse_providers() -> None: + checker_path = os.environ.get("COSMOS_CLUSTERPULSE_HEALTH_CHECK") + if not checker_path: + return + checker = Path(checker_path) + if not checker.is_absolute() or not checker.is_file() or not os.access(checker, os.X_OK): + raise RuntimeError(f"Missing executable ClusterPulse health checker: {checker}") + run_health_check(checker, process=False) + # The lock is per node and releases automatically when its process exits. + # It serializes initialization without storing readiness or creating a process group. + lock_path = Path(os.environ["SLURM_LOG_DIR"]) / f"clusterpulse-provider-init-{socket.gethostname()}.lock" + deadline = time.monotonic() + 60 + with lock_path.open("a") as lock: + while True: + try: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + break + except BlockingIOError: + if time.monotonic() >= deadline: + raise TimeoutError( + "Timed out waiting for node-local ClusterPulse provider initialization" + ) from None + time.sleep(0.1) + try: + emit_nvtx_probe() + run_health_check(checker, process=True) + finally: + fcntl.flock(lock, fcntl.LOCK_UN) + print(f"ClusterPulse providers verified: rank={os.environ['SLURM_PROCID']} pid={os.getpid()}", flush=True) diff --git a/cosmos_framework/utils/generator/cost_model/benchmark.py b/cosmos_framework/utils/generator/cost_model/benchmark.py index 379c5af5e..00f09d5e6 100644 --- a/cosmos_framework/utils/generator/cost_model/benchmark.py +++ b/cosmos_framework/utils/generator/cost_model/benchmark.py @@ -25,6 +25,11 @@ Sweeping inside one process is not a convenience: building the model and paying ``torch.compile`` warm-up costs minutes, while a timed point costs seconds, so re-launching per shape would spend almost all of its wall time on setup. +Automatic Python garbage collection is disabled during each point, including warm-up. +Generation-1 collection runs after each point, outside its timing and profile capture. +Raw-video release after preprocessing is always disabled so synthetic batches remain +available across warm-up and timed CP windows. This retains more GPU memory than +training with raw-video release enabled. The benchmark drives the model directly rather than through :class:`~cosmos_framework.trainer.ImaginaireTrainer`: it builds the model from a @@ -72,8 +77,11 @@ from __future__ import annotations import contextlib +import gc +import json import math import os +import socket import subprocess import time from dataclasses import dataclass, field @@ -85,9 +93,9 @@ import torch.distributed as dist import tyro -from cosmos_framework.data.generator.sequence_packing import SequencePlan -from cosmos_framework.utils import distributed, log, misc from cosmos_framework.utils.config import load_config +from cosmos_framework.utils.lazy_config import instantiate +from cosmos_framework.utils import distributed, log, misc from cosmos_framework.utils.generator.cost_model.estimator import ( PEAK_TFLOPS_BY_DEVICE, BucketMeasurement, @@ -98,13 +106,76 @@ flops_per_sample, ) from cosmos_framework.utils.generator.cost_model.spec import SampleSpec +from cosmos_framework.data.generator.sequence_packing import SequencePlan from cosmos_framework.utils.generator.method_timer import MethodTimer -from cosmos_framework.utils.lazy_config import instantiate CONFIG_PATH = "cosmos_framework/configs/base/config.py" DEFAULT_EXPERIMENT = "t2w_mot_exp306_006_qwen3_vl_8b_ratio_sample_balancing" +def configure_collective_debug(output_dir: str, steps: int) -> None: + """Configure opt-in diagnostics before NCCL or its flight recorder initializes.""" + recorder = getattr(torch._C, "_distributed_c10d", None) + for name in ("_dump_nccl_trace", "_reset_fr_recording_nccl"): + if not callable(getattr(recorder, name, None)): + raise RuntimeError(f"--debug-collectives requires this PyTorch build to provide {name}.") + for name in ("_enable_collectives_timing", "_wait_for_pending_works"): + if not callable(getattr(dist.ProcessGroup, name, None)): + raise RuntimeError(f"--debug-collectives requires ProcessGroup.{name}.") + if not hasattr(getattr(dist.distributed_c10d, "_world", None), "pg_map"): + raise RuntimeError("--debug-collectives requires the PyTorch process-group registry.") + + # A 94-layer CP32 step has over 2,500 collectives. Keep the entire timed + # point, not just the final 2,000 events in the launcher's default buffer. + buffer_size = max(65536, steps * 8192) + os.environ["TORCH_FR_BUFFER_SIZE"] = str(buffer_size) + os.environ["TORCH_NCCL_TRACE_BUFFER_SIZE"] = str(buffer_size) + debug_dir = Path(output_dir) / "collective_debug" + debug_dir.mkdir(parents=True, exist_ok=True) + os.environ["NCCL_DEBUG"] = "INFO" + os.environ["NCCL_DEBUG_SUBSYS"] = "INIT,ENV,NET,REG,ALLOC" + os.environ["NCCL_DEBUG_FILE"] = str(debug_dir / "nccl.%h.%p.log") + + +def dump_collective_debug( + output: Path, + groups: list[dist.ProcessGroup], + step_times: list[dict[str, int]], + timed_end_ns: int, +) -> None: + """Write all this rank's collectives after measurement and watchdog retirement. + + Flight Recorder's ``duration_ms`` is CUDA-event elapsed time. Its + ``time_discovered_started_ns`` and ``time_discovered_completed_ns`` are CPU + observation times, potentially delayed by the watchdog; they are not exact + GPU timestamps. Linux ``time_created_ns`` and the step markers both use + CLOCK_REALTIME. Cross-host clock offsets still need to be checked offline. + """ + # Synchronization at the end of time_bucket completes GPU work, but the + # watchdog still needs to retire it before duration_ms is present in dumps. + # This waits locally and releases the GIL; it adds no distributed barrier. + for group in groups: + group._wait_for_pending_works() + trace = torch._C._distributed_c10d._dump_nccl_trace( + includeCollectives=True, includeStackTraces=False, onlyActive=False + ) + output.parent.mkdir(parents=True, exist_ok=True) + output.with_suffix(".pickle").write_bytes(trace) + metadata = { + "rank": dist.get_rank(), + "world_size": dist.get_world_size(), + "hostname": socket.gethostname(), + "pid": os.getpid(), + "bucket": output.name, + "clock": "CLOCK_REALTIME nanoseconds; discovery timestamps are CPU observations", + "timed_end_ns": timed_end_ns, + "steps": step_times, + "buffer_size": int(os.environ["TORCH_FR_BUFFER_SIZE"]), + } + output.with_suffix(".json").write_text(json.dumps(metadata, indent=2) + "\n") + log.info(f"Wrote collective diagnostics: {output.with_suffix('.pickle')}") + + def _peak_tflops(device_name: str) -> float: """Dense BF16 peak for a CUDA device name, or 0.0 when it is not in the table.""" for known, peak in PEAK_TFLOPS_BY_DEVICE.items(): @@ -208,7 +279,7 @@ def build_model( Args: experiment: Registered Hydra experiment name. world_size: Ranks in this job. - extra_overrides: Additional OmegaConf overrides, applied last. + extra_overrides: Additional OmegaConf overrides. Raw-video release remains disabled. context_parallel_shard_degree: Context-parallel group size. Must divide ``world_size``; validated by the caller before this is reached. @@ -236,8 +307,10 @@ def build_model( "callbacks=[basic,optimization,job_monitor]", "model.config.ema.enabled=False", ] - # Applied last so an explicit override on the command line still wins. + # Explicit overrides win for configurable settings; raw-video retention is required below. overrides.extend(extra_overrides) + # Reuse the synthetic raw batch after warm-up and at every subsequent CP window. + overrides.append("model.config.release_cp_raw_batch_after_preprocessing=False") log.info(f"Loading {experiment} with overrides {overrides}") config = load_config(CONFIG_PATH, overrides) @@ -358,10 +431,15 @@ def time_bucket( step_optimizer: bool = True, vae_timer: MethodTimer | None = None, profiler: torch.profiler.profile | None = None, + collective_debug_output: Path | None = None, ) -> tuple[float, dict[str, Any], float, int]: - """Time ``steps`` steps (rounded up, see below) after ``warmup``, returning + """Time exactly ``steps`` steps after ``warmup``, returning ``(sec_per_step, last_output, vae_sec_per_step, steps_run)``. + Automatic Python GC is disabled for the entire point. Generation-1 collection, + matching production's collection level, runs afterward outside timing/profiling, + including on failure; the caller's previous automatic-GC state is restored. + Warm-up matters more than usual here: this experiment enables ``torch.compile``, so the first step of every new shape pays compilation. @@ -370,59 +448,97 @@ def time_bucket( a mean would report a step nobody experienced. ``vae_sec_per_step`` follows the same MAX convention when ``vae_timer`` is given; it is ``0.0`` otherwise. - ``profiler``, when given, is stepped once per TIMED iteration only, never during - warm-up. Its ``wait`` window (``--profile-wait-steps``) is documented as skipping + ``profiler``, when given, is entered after warm-up and stepped once per TIMED + iteration only. Its ``wait`` window (``--profile-wait-steps``) is documented as skipping steps on top of this bucket's own ``--warmup``; stepping it during warm-up too would run that wait window (and possibly its active window) out against compiling/warm-up steps instead, so the recorded trace would show compilation or warm-up noise rather than the steady-state timed steps it is meant to show. """ - for iteration in range(warmup): - run_step(model, optimizer, scheduler, batch, iteration, step_optimizer=step_optimizer) - - torch.cuda.synchronize() - if dist.is_initialized(): - dist.barrier() - - # OmniMoTModel only calls VAE encode on CP window slot 0 (see - # OmniMoTModel._get_training_inputs) -- once every cp_size calls to - # training_step, with the other cp_size-1 calls broadcasting a cached payload - # instead. Warm-up leaves the model's window slot wherever it happened to land, - # so without correcting it here the timed region starts mid-window and, unless - # ``steps`` is itself a multiple of cp_size, clips a fractional window at the - # far end too -- biasing vae_sec_per_step (and adding small-N noise to - # sec_per_step itself) by how far ``steps`` sits from a cp_size multiple. - # Forcing the slot to 0 and rounding steps up to the next whole multiple makes - # the timed region exactly ``steps // cp_size`` complete windows, each - # contributing exactly one encode -- the same cadence production trains at. - # Safe to poke the model's counter directly: this benchmark never runs - # ImaginaireTrainer (run_step calls the model directly, see its docstring), so - # there is no ContextParallelDataWindow tracking this rank's slot to desync - # from. - cp_size = getattr(getattr(model, "parallel_dims", None), "cp_size", 1) - if cp_size > 1: - model._cp_window_slot = 0 - steps = -(-steps // cp_size) * cp_size # ceil to the next multiple of cp_size - - if vae_timer is not None: - vae_timer.reset() - - start = time.perf_counter() - output_batch: dict[str, Any] = {} - for iteration in range(steps): - output_batch = run_step(model, optimizer, scheduler, batch, warmup + iteration, step_optimizer=step_optimizer) - if profiler is not None: - profiler.step() - torch.cuda.synchronize() - sec_per_step = (time.perf_counter() - start) / steps - - vae_sec_per_step = vae_timer.stop_and_read_ms() / 1000.0 / steps if vae_timer is not None else 0.0 - - if dist.is_initialized(): - worst = torch.tensor([sec_per_step, vae_sec_per_step], device="cuda") - dist.all_reduce(worst, op=dist.ReduceOp.MAX) - sec_per_step, vae_sec_per_step = float(worst[0].item()), float(worst[1].item()) - return sec_per_step, output_batch, vae_sec_per_step, steps + gc_enabled = gc.isenabled() + gc.disable() + try: + debug_groups: list[dist.ProcessGroup] = [] + debug_step_times: list[dict[str, int]] = [] + if collective_debug_output is not None: + debug_groups = [group for group in dist.distributed_c10d._world.pg_map if dist.get_backend(group) == "nccl"] + # Enable before warm-up so event allocation is warmed too. Earlier + # buckets retain their normal timing policy; there is no disable API. + for group in debug_groups: + group._enable_collectives_timing() + cp_size = getattr(getattr(model, "parallel_dims", None), "cp_size", 1) + if cp_size > 1: + # A partial prior window must not make this bucket warm up on the old shape. + model._cp_window_slot = 0 + model._cp_local_training_payload = None + for iteration in range(warmup): + run_step(model, optimizer, scheduler, batch, iteration, step_optimizer=step_optimizer) + + torch.cuda.synchronize() + if dist.is_initialized(): + dist.barrier() + + # OmniMoTModel only calls VAE encode on CP window slot 0 (see + # OmniMoTModel._get_training_inputs) -- once every cp_size calls to + # training_step, with the other cp_size-1 calls broadcasting a cached payload + # instead. Reset the window after warm-up so timing always starts at slot 0, + # independent of the warm-up count. The requested step count is never rounded: + # a partial final window overrepresents window-start overhead compared with + # production's complete windows. Measured VAE encoding is subtracted later, + # but other preparation/communication overhead remains in sec_per_step. + # Safe to poke the model's counter directly: this benchmark never runs + # ImaginaireTrainer (run_step calls the model directly, see its docstring), so + # there is no ContextParallelDataWindow tracking this rank's slot to desync + # from. + if cp_size > 1: + model._cp_window_slot = 0 + model._cp_local_training_payload = None + + if vae_timer is not None: + vae_timer.reset() + + if collective_debug_output is not None: + torch._C._distributed_c10d._reset_fr_recording_nccl() + timed_end_ns = 0 + with profiler if profiler is not None else contextlib.nullcontext(): + start = time.perf_counter() + output_batch: dict[str, Any] = {} + for iteration in range(steps): + if collective_debug_output is not None: + debug_step_times.append({"iteration": iteration, "cpu_start_ns": time.time_ns()}) + output_batch = run_step( + model, optimizer, scheduler, batch, warmup + iteration, step_optimizer=step_optimizer + ) + if collective_debug_output is not None: + debug_step_times[-1]["cpu_return_ns"] = time.time_ns() + if profiler is not None: + profiler.step() + torch.cuda.synchronize() + sec_per_step = (time.perf_counter() - start) / steps + if collective_debug_output is not None: + timed_end_ns = time.time_ns() + + if collective_debug_output is not None: + dump_collective_debug(collective_debug_output, debug_groups, debug_step_times, timed_end_ns) + vae_sec_per_step = vae_timer.stop_and_read_ms() / 1000.0 / steps if vae_timer is not None else 0.0 + + if dist.is_initialized(): + worst = torch.tensor([sec_per_step, vae_sec_per_step], device="cuda") + dist.all_reduce(worst, op=dist.ReduceOp.MAX) + sec_per_step, vae_sec_per_step = float(worst[0].item()), float(worst[1].item()) + if cp_size > 1: + # Release a partial window before the next bucket allocates its raw batch. + model._cp_window_slot = 0 + model._cp_local_training_payload = None + return sec_per_step, output_batch, vae_sec_per_step, steps + finally: + # Production uses generation 1. Collect between points so a rank cannot + # enter automatic GC mid-step and leave its peers waiting in collectives. + try: + gc.collect(1) + finally: + if gc_enabled: + gc.enable() def observed_spec( @@ -555,10 +671,9 @@ class BenchmarkConfig: Calibration.load_all.""" steps: int = 8 - """Timed steps per bucket. Under CP (``--context-parallel-shard-degree`` > 1), - ``time_bucket`` rounds this up to the next multiple of the CP degree so the timed - region covers whole encode windows -- see its docstring -- so the actual step count - a bucket ran at (``BucketMeasurement.steps``) may exceed this value.""" + """Exact number of timed steps per bucket, independent of CP size. Counts that + are not multiples of CP include a partial window and overrepresent window-start + preparation/communication overhead outside VAE encode.""" warmup: int = 3 """Untimed steps per bucket, which is what covers torch.compile.""" @@ -588,11 +703,11 @@ class BenchmarkConfig: """Wrap the timed steps in ``torch.profiler`` and write a trace per rank to ``/torch_profiles/rank``. Off by default: the profiler's own overhead lands on whichever bucket is profiled, so a calibration run should not combine - ``--profile`` with the points it means to fit against -- use it to diagnose a single - shape, not to run the sweep. Only the FIRST measured bucket is profiled (see - ``--profile-active-steps``); every later bucket runs exactly as it would without - ``--profile``, since a warm profiler with an exhausted schedule costs only a per-step - counter increment. Traces open in TensorBoard's PyTorch Profiler plugin or + ``--profile`` with the points it means to fit against. Only the selected bucket + is profiled (the first by default; see ``--profile-bucket`` and + ``--profile-active-steps``); all other buckets run without a profiler, preserving + the sweep history leading up to the selected shape. Traces open in TensorBoard's + PyTorch Profiler plugin or chrome://tracing, written gzip-compressed (see ``tensorboard_trace_handler``'s ``use_gzip``) since an uncompressed per-op CUDA trace across every rank is the largest thing this script writes and both consumers read ``.gz`` natively; the @@ -600,6 +715,14 @@ class BenchmarkConfig: ``tensorboard_trace_handler``'s own naming, not this script's, so it stays fixed regardless of compression.""" + debug_collectives: bool = False + """Capture collective CPU submissions and CUDA-event durations on ALL ranks for + all timed steps of the --profile-bucket point, without requiring --profile. Write Flight Recorder + dumps and step-clock metadata under collective_debug/rank, plus per-process + NCCL INFO logs. Recording adds overhead, so use this for diagnosis, not fitting + calibration coefficients. GPU start/end discovery timestamps are CPU observations, + not exact GPU timestamps. Off by default; requires the container's private Torch APIs.""" + profile_wait_steps: int = 1 """Untimed steps the profiler skips before it starts recording, on top of this bucket's own ``--warmup``. Lets the first step or two of the profiled window itself @@ -611,8 +734,9 @@ class BenchmarkConfig: per-op CUDA events is far more than needed to read off a kernel breakdown.""" profile_rank: int = 0 - """Which rank's ``torch.profiler`` context is active; every other rank runs the - profiled bucket with ``--profile`` effectively off. All ranks still execute the + """Which rank's ``torch.profiler`` context is active unless --profile-ranks is given; + every unselected rank runs the profiled bucket with ``--profile`` effectively off. + All ranks still execute the same steps and the same collectives -- profiling is local CPU/GPU instrumentation, not a collective itself -- so this only cuts the per-op recording overhead (and trace file size) on the ranks that would otherwise write a trace nobody asked for. @@ -621,6 +745,14 @@ class BenchmarkConfig: ``sec_per_step``; see ``--profile``'s docstring for why a profiled run is a diagnostic, not a calibration point, regardless of how many ranks are traced.""" + profile_ranks: list[int] | None = None + """Global ranks to profile simultaneously. Overrides --profile-rank when supplied.""" + + profile_bucket: str | None = None + """Exact bucket label for profiling or collective diagnostics, e.g. video_720_173f + (with @txt4096 when sweeping captions). + Defaults to the first bucket; preceding buckets run normally with no profiler attached.""" + log_recompiles: bool = True """Enable dynamo's TORCH_LOGS=recompiles logging for the run. On by default: a recompile mid-sweep is exactly the kind of thing that silently inflates sec_per_step @@ -680,12 +812,20 @@ def __post_init__(self) -> None: raise ValueError(f"--profile-active-steps must be positive, got {self.profile_active_steps}.") if self.profile_rank < 0: raise ValueError(f"--profile-rank cannot be negative, got {self.profile_rank}.") + if self.profile_ranks is not None: + if not self.profile_ranks or any(rank < 0 for rank in self.profile_ranks): + raise ValueError("--profile-ranks needs at least one nonnegative rank.") + self.profile_ranks = list(dict.fromkeys(self.profile_ranks)) + if self.profile and self.profile_wait_steps + self.profile_active_steps > self.steps: + raise ValueError("--steps must cover --profile-wait-steps plus --profile-active-steps.") def main(args: BenchmarkConfig) -> None: if args.log_recompiles: torch._logging.set_logs(recompiles=True) + if args.debug_collectives: + configure_collective_debug(args.output_dir, args.steps) # NCCL reads this when the communicator is created. Match the training # launcher's NVLS policy for this CLI, while allowing an explicit shell # override. Importing the reusable benchmark helpers must not change the @@ -704,8 +844,9 @@ def main(args: BenchmarkConfig) -> None: f"world_size={world_size}: CP ranks overlay onto dp rank slots in equal-sized groups, so the " "degree must evenly partition the launched ranks." ) - if args.profile_rank >= world_size: - raise ValueError(f"--profile-rank={args.profile_rank} is out of range for world_size={world_size}.") + profile_ranks = args.profile_ranks if args.profile_ranks is not None else [args.profile_rank] + if any(profile_rank >= world_size for profile_rank in profile_ranks): + raise ValueError(f"Profile ranks {profile_ranks} are out of range for world_size={world_size}.") frame_counts = args.frames text_token_counts = args.text_tokens @@ -780,6 +921,10 @@ def bucket_label(spec: SampleSpec) -> str: label += f"@txt{spec.text_tokens}" return label + profile_bucket = args.profile_bucket if args.profile_bucket is not None else bucket_label(points[0][0]) + if (args.profile or args.debug_collectives) and profile_bucket not in {bucket_label(spec) for spec, _ in points}: + raise ValueError(f"--profile-bucket={profile_bucket!r} is not in the requested sweep.") + log.info( f"Sweeping {len(points)} point(s) on {world_size} rank(s) at batch sizes filling a " f"{args.tokens_per_step:,}-token step per rank, and caption lengths {text_token_counts}" @@ -792,6 +937,7 @@ def bucket_label(spec: SampleSpec) -> str: # copy is ever built, hence nothing to state about it here. carried = "with the parameter update" if args.optimizer else "forward and backward only" log.info(f"Timed step: {carried}, no EMA copy, VAE AOTI-compiled") + log.info("Python GC: automatic collection disabled during each point; manual generation-1 collection afterward") model, optimizer, scheduler = build_model( args.experiment, world_size, overrides, context_parallel_shard_degree=args.context_parallel_shard_degree @@ -807,6 +953,12 @@ def bucket_label(spec: SampleSpec) -> str: # ParallelDims actually built the cp mesh at, read the same way any other # consumer of model.parallel_dims would. cp_size = model.parallel_dims.cp_size + if args.steps % cp_size: + log.warning( + f"Timing exactly {args.steps} steps at CP{cp_size}, including a partial CP window. " + "Measured VAE encode time is subtracted, but other window-start preparation/communication " + "overhead is overrepresented compared with a full CP window." + ) if cp_size > 1: log.info( f"Context parallelism: cp={cp_size}. Ranks within a cp group of {cp_size} split ONE sample's packed " @@ -833,125 +985,129 @@ def bucket_label(spec: SampleSpec) -> str: vae_timer = MethodTimer(model, "encode") profile_dir = str(Path(args.output_dir) / "torch_profiles") - profile_this_rank = args.profile and rank == args.profile_rank - profiler_cm = ( - torch.profiler.profile( - activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], - schedule=torch.profiler.schedule( - wait=args.profile_wait_steps, warmup=0, active=args.profile_active_steps, repeat=1 - ), - on_trace_ready=torch.profiler.tensorboard_trace_handler(f"{profile_dir}/rank{rank}", use_gzip=True), - ) - if profile_this_rank - else contextlib.nullcontext() - ) + profile_this_rank = args.profile and rank in profile_ranks if profile_this_rank: log.info( - f"Profiling the first measured bucket on rank {rank} only: {args.profile_wait_steps} step(s) " - f"skipped, then {args.profile_active_steps} recorded, trace written to {profile_dir}. Every other " - "rank runs the same steps and collectives without the profiler attached." + f"Profiling bucket {profile_bucket} on ranks {profile_ranks}: {args.profile_wait_steps} step(s) " + f"skipped, then {args.profile_active_steps} recorded, trace written to {profile_dir}/rank{rank}. " + "Unselected ranks and buckets run without the profiler attached." ) - elif args.profile: - log.info(f"--profile is on but this is rank {rank}; only rank {args.profile_rank} attaches a profiler.") measurements: dict[str, BucketMeasurement] = {} # Smallest packed length that ran out of memory, per resolution. Longer # clips at that resolution cannot fit either, so they are never attempted. oom_floor: dict[str, int] = {} skipped: dict[str, str] = {} - with profiler_cm as profiler: - for point_index, (spec, batch_size) in enumerate(points): - label = bucket_label(spec) - floor = oom_floor.get(spec.resolution) - if floor is not None and spec.total_tokens >= floor: - reason = f"skipped: {spec.resolution}p ran out of memory at {floor} tokens" - log.warning(f"{label}: {reason}") - skipped[label] = reason - continue - - height, width = spec.pixel_shape - modality = spec.modality - failed = False - sec_per_step = 0.0 - vae_sec_per_step = 0.0 - steps_run = args.steps - output_batch: dict[str, Any] = {} - try: - batch = synthetic_batch(modality, height, width, spec.num_pixel_frames, batch_size, spec.text_tokens) - torch.cuda.reset_peak_memory_stats() - sec_per_step, output_batch, vae_sec_per_step, steps_run = time_bucket( - model, - optimizer, - scheduler, - batch, - args.steps, - args.warmup, - step_optimizer=args.optimizer, - vae_timer=vae_timer, - # Only the first bucket carries the profiler's overhead -- see - # BenchmarkConfig.profile. - profiler=profiler if point_index == 0 else None, + for spec, batch_size in points: + label = bucket_label(spec) + floor = oom_floor.get(spec.resolution) + if floor is not None and spec.total_tokens >= floor: + reason = f"skipped: {spec.resolution}p ran out of memory at {floor} tokens" + log.warning(f"{label}: {reason}") + skipped[label] = reason + continue + + height, width = spec.pixel_shape + modality = spec.modality + failed = False + sec_per_step = 0.0 + vae_sec_per_step = 0.0 + steps_run = args.steps + output_batch: dict[str, Any] = {} + try: + batch = synthetic_batch(modality, height, width, spec.num_pixel_frames, batch_size, spec.text_tokens) + torch.cuda.reset_peak_memory_stats() + # Construct and enter a profiler only for the selected bucket. Earlier + # buckets retain their training history without any profiling overhead. + profiler = ( + torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], + schedule=torch.profiler.schedule( + wait=args.profile_wait_steps, warmup=0, active=args.profile_active_steps, repeat=1 + ), + on_trace_ready=torch.profiler.tensorboard_trace_handler(f"{profile_dir}/rank{rank}", use_gzip=True), ) - except torch.OutOfMemoryError: - failed = True - log.warning(f"{label}: out of memory at {spec.total_tokens} tokens x bs{batch_size}") - - # Ranks vote so they abandon the point together; a one-sided skip would - # desync the next point's collectives. - if any_rank(failed): - floor = min(spec.total_tokens, oom_floor.get(spec.resolution, spec.total_tokens)) - oom_floor[spec.resolution] = floor - skipped[label] = f"out of memory at {spec.total_tokens} tokens x bs{batch_size}" - continue - - spec = observed_spec( - output_batch, batch_size, modality, spec.resolution, args.aspect_ratio, spec.num_pixel_frames - ) - note = check_vision_tokens(spec, output_batch, batch_size) - - # Subtract this run's own MEASURED VAE-encode time (via vae_timer / CUDA events) - # rather than a predicted one, so the fitted cost model prices transformer compute - # only. Measuring beats predicting here because the step is actually running -- - # predicted_encode_seconds exists for decisions that have to be made before a step - # runs (VAE load balancing), not for after-the-fact accounting where the real - # number is directly available. - sec_per_step = max(sec_per_step - vae_sec_per_step, 0.0) - - # A cp group of cp_size ranks cooperates on ONE logical batch -- each rank - # narrows its own locally built sequence down to its 1/cp_size shard rather - # than contributing a distinct sample (see context_parallel_utils.py). So - # only one cp group's worth of batch_size samples is real work per group, - # and there are world_size / cp_size such groups in the step. - samples_per_step = batch_size * world_size // cp_size - gpu_sec_per_step = sec_per_step * world_size - sample_flops = flops_per_sample(descriptor, spec, flags) - measurements[label] = BucketMeasurement( - name=label, - spec=spec, - steps=steps_run, - samples_per_step=samples_per_step, - sec_per_step=sec_per_step, - gpu_sec_per_sample=gpu_sec_per_step / samples_per_step, - flops_per_sample=sample_flops, - achieved_tflops_per_gpu=sample_flops * samples_per_step / gpu_sec_per_step / 1e12, - peak_mem_gb=torch.cuda.max_memory_allocated() / 2**30, - world_size=world_size, - note=note, + if profile_this_rank and label == profile_bucket + else None ) - # Both units, because the two numbers that matter live in different ones: the - # token target is per rank, while samples_per_step is global and is what the fit - # scales its per-sample columns by. Printing the per-sample token count against - # the global sample count invites multiplying them into a figure that describes - # neither a rank nor the target. - cp_suffix = f", cp={cp_size}" if cp_size > 1 else "" - log.info( - f"{label}: {sec_per_step * 1000:.1f} ms/step (transformer-only, VAE " - f"{vae_sec_per_step * 1000:.1f} ms/step measured & subtracted), " - f"{spec.total_tokens} tokens x {batch_size} samples/rank = {batch_size * spec.total_tokens:,} " - f"tokens/rank ({samples_per_step} distinct samples/step{cp_suffix}), " - f"{measurements[label].achieved_tflops_per_gpu:.1f} TFLOP/s/GPU, " - f"{measurements[label].peak_mem_gb:.1f} GiB peak" + sec_per_step, output_batch, vae_sec_per_step, steps_run = time_bucket( + model, + optimizer, + scheduler, + batch, + args.steps, + args.warmup, + step_optimizer=args.optimizer, + vae_timer=vae_timer, + # Only the selected bucket carries the profiler's overhead -- see + # BenchmarkConfig.profile. + profiler=profiler, + collective_debug_output=( + Path(args.output_dir) / "collective_debug" / f"rank{rank}" / label + if args.debug_collectives and label == profile_bucket + else None + ), ) + except torch.OutOfMemoryError: + failed = True + log.warning(f"{label}: out of memory at {spec.total_tokens} tokens x bs{batch_size}") + + # Ranks vote so they abandon the point together; a one-sided skip would + # desync the next point's collectives. + if any_rank(failed): + floor = min(spec.total_tokens, oom_floor.get(spec.resolution, spec.total_tokens)) + oom_floor[spec.resolution] = floor + skipped[label] = f"out of memory at {spec.total_tokens} tokens x bs{batch_size}" + continue + + spec = observed_spec( + output_batch, batch_size, modality, spec.resolution, args.aspect_ratio, spec.num_pixel_frames + ) + note = check_vision_tokens(spec, output_batch, batch_size) + + # Subtract this run's own MEASURED VAE-encode time (via vae_timer / CUDA events) + # rather than a predicted one, so the fitted cost model prices transformer compute + # only. Measuring beats predicting here because the step is actually running -- + # predicted_encode_seconds exists for decisions that have to be made before a step + # runs (VAE load balancing), not for after-the-fact accounting where the real + # number is directly available. + sec_per_step = max(sec_per_step - vae_sec_per_step, 0.0) + + # A cp group of cp_size ranks cooperates on ONE logical batch -- each rank + # narrows its own locally built sequence down to its 1/cp_size shard rather + # than contributing a distinct sample (see context_parallel_utils.py). So + # only one cp group's worth of batch_size samples is real work per group, + # and there are world_size / cp_size such groups in the step. + samples_per_step = batch_size * world_size // cp_size + gpu_sec_per_step = sec_per_step * world_size + sample_flops = flops_per_sample(descriptor, spec, flags) + measurements[label] = BucketMeasurement( + name=label, + spec=spec, + steps=steps_run, + samples_per_step=samples_per_step, + sec_per_step=sec_per_step, + gpu_sec_per_sample=gpu_sec_per_step / samples_per_step, + flops_per_sample=sample_flops, + achieved_tflops_per_gpu=sample_flops * samples_per_step / gpu_sec_per_step / 1e12, + peak_mem_gb=torch.cuda.max_memory_allocated() / 2**30, + world_size=world_size, + note=note, + ) + # Both units, because the two numbers that matter live in different ones: the + # token target is per rank, while samples_per_step is global and is what the fit + # scales its per-sample columns by. Printing the per-sample token count against + # the global sample count invites multiplying them into a figure that describes + # neither a rank nor the target. + cp_suffix = f", cp={cp_size}" if cp_size > 1 else "" + log.info( + f"{label}: {sec_per_step * 1000:.1f} ms/step (transformer-only, VAE " + f"{vae_sec_per_step * 1000:.1f} ms/step measured & subtracted), " + f"{spec.total_tokens} tokens x {batch_size} samples/rank = {batch_size * spec.total_tokens:,} " + f"tokens/rank ({samples_per_step} distinct samples/step{cp_suffix}), " + f"{measurements[label].achieved_tflops_per_gpu:.1f} TFLOP/s/GPU, " + f"{measurements[label].peak_mem_gb:.1f} GiB peak" + ) if not measurements: raise RuntimeError("Every sweep point failed or was skipped; nothing to calibrate against.") @@ -978,6 +1134,8 @@ def bucket_label(spec: SampleSpec) -> str: "overrides": overrides, "steps": args.steps, "warmup": args.warmup, + "gc_policy": "manual_generation_1_after_each_point", + "debug_collectives": args.debug_collectives, "swept_resolutions": resolutions, "swept_frames": frame_counts, # The batch sizes the target worked out to, so a calibration says how each diff --git a/cosmos_framework/utils/generator/deterministic_rng.py b/cosmos_framework/utils/generator/deterministic_rng.py index b02504e28..e75f6d6f6 100644 --- a/cosmos_framework/utils/generator/deterministic_rng.py +++ b/cosmos_framework/utils/generator/deterministic_rng.py @@ -31,16 +31,29 @@ def sample_seed(data: dict[str, Any]) -> int: class DeterministicAugmentor(Augmentor): - def __init__(self, augmentor: Augmentor) -> None: - self.augmentor = augmentor + def __init__(self, augmentor: Augmentor, *, fixed_sample_epoch: int | None = None) -> None: + """Optionally pin the augmented view for a measured offline packing catalog.""" + if fixed_sample_epoch is not None and (type(fixed_sample_epoch) is not int or fixed_sample_epoch < 0): + raise ValueError("fixed_sample_epoch must be a non-negative integer or None.") + self.augmentor: Augmentor = augmentor + self.fixed_sample_epoch: int | None = fixed_sample_epoch def __call__(self, data: dict[str, Any]) -> dict[str, Any]: + if self.fixed_sample_epoch is not None: + if not isinstance(data, dict) or not isinstance(data.get("__key__"), str) or not data["__key__"]: + raise ValueError("A fixed augmented view requires a non-empty sample key.") + url = data.get("__url__", "") + root = getattr(url, "root", url) + if not isinstance(root, str) or not root: + raise ValueError("A fixed augmented view requires a non-empty source root.") if not isinstance(data, dict) or "__key__" not in data: return self.augmentor(data) python_state = random.getstate() numpy_state = np.random.get_state() torch_state = torch.get_rng_state() # [N_state] - seed = sample_seed(data) + # Only the RNG identity is pinned; preserve the caller's epoch metadata. + seed_data = data if self.fixed_sample_epoch is None else {**data, "sample_epoch": self.fixed_sample_epoch} + seed = sample_seed(seed_data) random.seed(seed) np.random.seed(seed % (2**32)) torch_generator = torch.Generator(device="cpu") diff --git a/cosmos_framework/utils/generator/fused_adam.py b/cosmos_framework/utils/generator/fused_adam.py index 984b27a6b..9c70b4941 100644 --- a/cosmos_framework/utils/generator/fused_adam.py +++ b/cosmos_framework/utils/generator/fused_adam.py @@ -231,6 +231,12 @@ def step(self, closure=None, grads=None, output_params=None, scale=None, grad_no # Exponential moving average of squared gradient values state["exp_avg_sq"] = torch.zeros_like(p).float() + # Keep empty local shards in optimizer state for checkpoint resharding, + # but omit them from TE's kernel lists. Mixing empty and nonempty + # tensors in the fused kernel can cause a CUDA memory access error. + if get_local_tensor_if_DTensor(p).numel() == 0: + continue + if p.dtype == torch.float16: if self.master_weights: p_16_master.append(get_local_tensor_if_DTensor(p_master)) diff --git a/cosmos_framework/utils/generator/hf_attention_cosmos.py b/cosmos_framework/utils/generator/hf_attention_cosmos.py index fe36f78b1..d5429975b 100644 --- a/cosmos_framework/utils/generator/hf_attention_cosmos.py +++ b/cosmos_framework/utils/generator/hf_attention_cosmos.py @@ -56,9 +56,10 @@ from cosmos_framework.model.attention import attention as imag_attention from cosmos_framework.model.attention.masks import CausalType +from cosmos_framework.model.generator.context_parallel import sequence_all_to_all -def hf_attention_cosmos( +def _hf_attention_cosmos_local( module: torch.nn.Module, query: torch.Tensor, # [B, num_heads, N, head_dim] (BHSD) key: torch.Tensor, # [B, num_kv_heads, N_kv, head_dim] @@ -194,6 +195,8 @@ def hf_attention_cosmos( cumulative_seqlen_KV=cu_seq_lens_k, max_seqlen_Q=max_length_q, max_seqlen_KV=max_length_k, + # Backend selection has its own flag; the model/CLI sets the Torch mode. + deterministic=torch.are_deterministic_algorithms_enabled(), ) # imag_attention returns a bare Tensor unless return_lse=True (which we don't pass); # cast narrows the Tensor | tuple[Tensor, Tensor] return type. @@ -204,3 +207,105 @@ def hf_attention_cosmos( # out is BSHD [B, N, num_heads, head_dim_v] — matches HF's expected return shape. return out, None + + +def hf_attention_cosmos( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + dropout: float = 0.0, + scaling: float | None = None, + sliding_window: int | None = None, + cu_seq_lens_q: torch.Tensor | None = None, + cu_seq_lens_k: torch.Tensor | None = None, + max_length_q: int | None = None, + max_length_k: int | None = None, + decoder_attention_chunks: tuple[tuple[int, int, int, int, int], ...] | None = None, + **kwargs: Any, +) -> tuple[torch.Tensor, None]: + """Exchange sequence/head shards before dense or globally segmented attention. + + Packed boundaries remain global: each head owner sees the complete pack and + applies the same isolated per-conversation attention as CP1. Vision modules + have no CP group and retain their original attention path. + """ + cp_group = getattr(module, "_cosmos_cp_group", None) + if cp_group is not None: + if not getattr(module, "is_causal", False): + raise ValueError("CP requires causal decoder self-attention") + query, key, value = ( + sequence_all_to_all(tensor, cp_group, scatter_dim=1, gather_dim=2) for tensor in (query, key, value) + ) + if decoder_attention_chunks is not None: + if ( + not getattr(module, "is_causal", False) + or cu_seq_lens_q is None + or cu_seq_lens_k is not cu_seq_lens_q + or max_length_q != max_length_k + or query.shape[0] != 1 + or key.shape[0] != 1 + or value.shape[0] != 1 + or not (query.shape[2] == key.shape[2] == value.shape[2]) + ): + raise ValueError("Decoder attention chunks require B=1 causal self-attention with shared boundaries") + if ( + not decoder_attention_chunks + or decoder_attention_chunks[-1][1] != cu_seq_lens_q.numel() - 1 + or decoder_attention_chunks[-1][3] != query.shape[2] + ): + raise ValueError("Decoder attention chunks must cover the complete packed row") + end_sequence = end_token = 0 + outputs: list[torch.Tensor] = [] + for first, end, token_start, token_end, max_length in decoder_attention_chunks: + if ( + first != end_sequence + or token_start != end_token + or end <= first + or token_end <= token_start + or max_length <= 0 + or end >= cu_seq_lens_q.numel() + or token_end > query.shape[2] + ): + raise ValueError("Decoder attention chunks must be contiguous, positive and ordered") + boundaries = cu_seq_lens_q[first : end + 1] - token_start + output, _ = _hf_attention_cosmos_local( + module, + query[:, :, token_start:token_end], + key[:, :, token_start:token_end], + value[:, :, token_start:token_end], + attention_mask, + dropout=dropout, + scaling=scaling, + sliding_window=sliding_window, + cu_seq_lens_q=boundaries, + cu_seq_lens_k=boundaries, + max_length_q=max_length, + max_length_k=max_length, + **kwargs, + ) + outputs.append(output) + end_sequence, end_token = end, token_end + if end_sequence != cu_seq_lens_q.numel() - 1 or end_token != query.shape[2]: + raise ValueError("Decoder attention chunks must cover all sequences and tokens") + output = torch.cat(outputs, dim=1) if len(outputs) > 1 else outputs[0] + else: + output, _ = _hf_attention_cosmos_local( + module, + query, + key, + value, + attention_mask, + dropout=dropout, + scaling=scaling, + sliding_window=sliding_window, + cu_seq_lens_q=cu_seq_lens_q, + cu_seq_lens_k=cu_seq_lens_k, + max_length_q=max_length_q, + max_length_k=max_length_k, + **kwargs, + ) + if cp_group is not None: + output = sequence_all_to_all(output, cp_group, scatter_dim=1, gather_dim=2) + return output, None diff --git a/cosmos_framework/utils/generator/monkey_patch.py b/cosmos_framework/utils/generator/monkey_patch.py index bcbedfd0c..385c6022f 100644 --- a/cosmos_framework/utils/generator/monkey_patch.py +++ b/cosmos_framework/utils/generator/monkey_patch.py @@ -23,6 +23,8 @@ from transformers.utils.import_utils import is_torchdynamo_compiling from cosmos_framework.utils import log +from cosmos_framework.model.generator.context_parallel import shard_decoder_inputs +from cosmos_framework.utils.generator.attention_chunks import plan_varlen_attention_chunks _EXPECTED_TRANSFORMERS_VERSION_PREFIX = "4.57." @@ -91,9 +93,9 @@ def patched_forward( self, input_ids: torch.LongTensor = None, attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, + position_ids: torch.Tensor | None = None, past_key_values: Cache | None = None, - inputs_embeds: torch.FloatTensor | None = None, + inputs_embeds: torch.Tensor | None = None, pixel_values: torch.Tensor | None = None, pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, @@ -258,6 +260,36 @@ def patched_forward( position_ids = position_ids.add(delta) position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + cp_group = getattr(self, "_cosmos_multimodal_cp_group", None) + if cp_group is not None and (has_image or has_video): + if not all( + isinstance(value, torch.Tensor) + for value in (inputs_embeds, position_ids, attention_mask, visual_pos_masks) + ): + raise ValueError("Multimodal CP requires complete tensor embeddings, positions and masks") + assert isinstance(inputs_embeds, torch.Tensor) and isinstance(position_ids, torch.Tensor) + assert isinstance(attention_mask, torch.Tensor) and isinstance(visual_pos_masks, torch.Tensor) + if deepstack_visual_embeds is None: + raise ValueError("Multimodal CP requires the complete DeepStack features") + # Vision encoding, media insertion and native mRoPE all use the complete sequence. + local = shard_decoder_inputs( + { + "inputs_embeds": inputs_embeds, + "position_ids": position_ids, + "attention_mask": attention_mask, + "visual_pos_masks": visual_pos_masks, + "deepstack_visual_embeds": deepstack_visual_embeds, + }, + cp_group, + position_axes=3, + clone_embeddings=True, + ) + inputs_embeds = local["inputs_embeds"] + position_ids = local["position_ids"] + attention_mask = local["attention_mask"] + visual_pos_masks = local["visual_pos_masks"] + deepstack_visual_embeds = local["deepstack_visual_embeds"] + outputs = self.language_model( input_ids=None, position_ids=position_ids, @@ -589,9 +621,25 @@ def patched_moe_block_forward(self, hidden_states: torch.Tensor) -> tuple[torch. return len(blocks) -def patch_qwen3_vl_vision_varlen_attention(model: torch.nn.Module) -> int: +def _vision_attention_chunks( + grid_thw: list[list[int]], max_padded_tokens: int +) -> tuple[tuple[int, int, int, int, int], ...]: + """Group complete frame sequences while bounding batch_size * max_sequence_length. + + Entries are (first_sequence, end_sequence, first_token, end_token, max_length). + No frame is split or reordered. A single frame longer than the budget stays intact. + """ + if max_padded_tokens <= 0 or not grid_thw or any(min(grid) <= 0 for grid in grid_thw): + raise ValueError("Vision grids and the padded-token budget must be positive") + lengths = [height * width for time, height, width in grid_thw for _ in range(time)] + return plan_varlen_attention_chunks(lengths, max_padded_tokens) + + +def patch_qwen3_vl_vision_varlen_attention(model: torch.nn.Module, *, max_padded_tokens: int = 0) -> int: """Monkey-patch every Qwen3-VL vision tower **instance** under ``model`` so each block runs - ONE varlen attention call over the packed patches instead of one call per image/frame. + varlen attention over packed patches instead of one call per image/frame. By default + there is one call; an optional padded-token budget bounds backward workspaces when + high-resolution images are packed together with many small video frames. HF's ``Qwen3VL{,Moe}VisionAttention.forward`` reserves its varlen path for ``_attn_implementation == "flash_attention_2"``. Every other implementation — including @@ -619,7 +667,8 @@ def patch_qwen3_vl_vision_varlen_attention(model: torch.nn.Module) -> int: sdpa/eager overrides that choice silently. ``hf_model`` gates the call on it. Patched in place, per vision tower: - - Each ``Qwen3VL{,Moe}VisionAttention.forward`` becomes the FA2-shaped single call. + - Each ``Qwen3VL{,Moe}VisionAttention.forward`` becomes a packed varlen call, + or a bounded group of calls when ``max_padded_tokens`` is set. - The tower's own ``forward`` gains a wrapper that computes ``max_seqlen`` — an int, as the kernels need a host value to size their scheduling — and passes it down through the ``**kwargs`` that the tower already threads into every block. That leaves ONE @@ -635,6 +684,8 @@ def patch_qwen3_vl_vision_varlen_attention(model: torch.nn.Module) -> int: Args: model: Module containing the vision tower(s), i.e. the outer ``Qwen3VL{,Moe}ForConditionalGeneration`` (``HFModel.model``). + max_padded_tokens: Optional bound on sequence count times maximum frame length + per attention call. Zero preserves the original single-call path. Returns: Number of patched attention modules. @@ -649,6 +700,9 @@ def patch_qwen3_vl_vision_varlen_attention(model: torch.nn.Module) -> int: # backends at module scope, and backbones that never reach this patch must not pay for that. from cosmos_framework.utils.generator.hf_attention_cosmos import hf_attention_cosmos + if max_padded_tokens < 0: + raise ValueError("Vision attention padded-token budget must be nonnegative") + towers = [module for module in model.modules() if isinstance(module, (Qwen3VLVisionModel, Qwen3VLMoeVisionModel))] if not towers: raise ValueError( @@ -674,6 +728,7 @@ def patched_vision_attention_forward( rotary_pos_emb: torch.Tensor | None = None, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, max_seqlen: int | None = None, + vision_attention_chunks: tuple[tuple[int, int, int, int, int], ...] | None = None, **kwargs, ) -> torch.Tensor: if max_seqlen is None: @@ -709,20 +764,43 @@ def patched_vision_attention_forward( # the other implementations in that registry swallow unrecognized kwargs, so a tower # whose _attn_implementation changed after patching would drop them on the floor and # attend ACROSS image boundaries — wrong features, no error. - attn_output, _ = hf_attention_cosmos( - self, - query_states, - key_states, - value_states, - attention_mask=None, - scaling=self.scaling, - dropout=0.0 if not self.training else self.attention_dropout, - cu_seq_lens_q=cu_seqlens, - cu_seq_lens_k=cu_seqlens, - max_length_q=max_seqlen, - max_length_k=max_seqlen, - is_causal=False, - ) # [1,N,num_heads,head_dim] + if vision_attention_chunks is None: + attn_output, _ = hf_attention_cosmos( + self, + query_states, + key_states, + value_states, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + cu_seq_lens_q=cu_seqlens, + cu_seq_lens_k=cu_seqlens, + max_length_q=max_seqlen, + max_length_k=max_seqlen, + is_causal=False, + ) # [1,N,num_heads,head_dim] + else: + outputs: list[torch.Tensor] = [] + for first, end, token_start, token_end, chunk_max in vision_attention_chunks: + # All heads for each complete frame stay together. FSDP still sees + # one vision/module forward; only independent attention calls split. + chunk_cu = cu_seqlens[first : end + 1] - token_start + output, _ = hf_attention_cosmos( + self, + query_states[:, :, token_start:token_end], + key_states[:, :, token_start:token_end], + value_states[:, :, token_start:token_end], + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + cu_seq_lens_q=chunk_cu, + cu_seq_lens_k=chunk_cu, + max_length_q=chunk_max, + max_length_k=chunk_max, + is_causal=False, + ) + outputs.append(output) + attn_output = torch.cat(outputs, dim=1) # [1,N,num_heads,head_dim] attn_output = attn_output.reshape(seq_length, -1).contiguous() # [N,hidden_size] attn_output = self.proj(attn_output) # [N,hidden_size] @@ -746,7 +824,14 @@ def patched_vision_tower_forward( # sequences of the same length), so the longest sequence is max(h*w) — the same # value HF's FA2 branch derives from cu_seqlens, read here once for the whole # tower. int() is the device-to-host sync; it sits outside the compiled blocks. - kwargs.setdefault("max_seqlen", int((grid_thw[:, 1] * grid_thw[:, 2]).max())) + if max_padded_tokens: + # Plan once outside the blocks and reuse during activation recomputation. + # This replaces the existing host read; it adds none inside attention. + chunks = _vision_attention_chunks(grid_thw.tolist(), max_padded_tokens) + kwargs.setdefault("max_seqlen", max(chunk[-1] for chunk in chunks)) + kwargs["vision_attention_chunks"] = chunks if len(chunks) > 1 else None + else: + kwargs.setdefault("max_seqlen", int((grid_thw[:, 1] * grid_thw[:, 2]).max())) return tower_forward(hidden_states, grid_thw, **kwargs) return patched_vision_tower_forward diff --git a/cosmos_framework/utils/generator/parallelism.py b/cosmos_framework/utils/generator/parallelism.py index 8b3c64aa3..2893180b9 100644 --- a/cosmos_framework/utils/generator/parallelism.py +++ b/cosmos_framework/utils/generator/parallelism.py @@ -45,8 +45,6 @@ from cosmos_framework.utils import log -_MAX_CP = 64 - @dataclass class ParallelDims: @@ -64,7 +62,9 @@ class ParallelDims: consume the dp budget). dp_replicate: HSDP replicate size. Pass ``-1`` to auto-infer to ``world_size // dp_shard``. - cp: Context parallel size in ``[1, _MAX_CP]``. Overlay axis. + cp: Positive context parallel size. Overlay axis; ``cp * cfgp`` + must divide ``world_size``. The attention implementation + validates model-specific constraints such as query heads. cfgp: CFG parallel size in ``(1, 2)``. Overlay axis, inference-only (rejected at construction time unless ``enable_inference_mode`` is True). Also is @@ -107,8 +107,8 @@ def _validate(self) -> None: # --- overlay range checks (run before division below) --- if self.fsdp_cpu_offload and not self.enable_inference_mode: raise ValueError("fsdp_cpu_offload is supported only in inference mode") - if self.cp < 1 or self.cp > _MAX_CP: - raise ValueError(f"CP (Context Parallelism) must be in [1, {_MAX_CP}]. got {self.cp}") + if self.cp < 1: + raise ValueError(f"CP (Context Parallelism) must be positive. got {self.cp}") if self.cfgp not in (1, 2): raise ValueError(f"CFGP (CFG Parallelism) must be 1 or 2. got {self.cfgp}") if not self.enable_inference_mode and self.cfgp > 1: @@ -207,6 +207,8 @@ def build_meshes(self, device_type: str = "cuda") -> None: dims=[self.dp_replicate, self.dp_shard], names=["dp_replicate", "dp_shard"], ) + # Register a flattened DP mesh so DTensor can fuse reductions across both DP axes. + self._meshes["dp"]._flatten(mesh_dim_name="dp_flat") # Pure FSDP uses a 1-D mesh, including the singleton CPU-offload # case. Passing the full 2-D (1, N) mesh selects FSDP2's HSDP path. if self.dp_shard_enabled or not self.dp_replicate_enabled: diff --git a/cosmos_framework/utils/generator/reasoner/constant.py b/cosmos_framework/utils/generator/reasoner/constant.py index ed6119bda..001882f0a 100644 --- a/cosmos_framework/utils/generator/reasoner/constant.py +++ b/cosmos_framework/utils/generator/reasoner/constant.py @@ -21,3 +21,7 @@ PROCESSOR_KEYS_TO_ADD_EAGLE: list[str] = ["input_ids", "attention_mask", "pixel_values", "image_sizes"] PROCESSOR_KEYS_TO_ADD: list[str] = list(set(PROCESSOR_KEYS_TO_ADD_QWEN + PROCESSOR_KEYS_TO_ADD_EAGLE)) + +# Conversation roles whose content can carry images, videos and audio: user turns, and tool results (the +# Qwen3-VL chat template renders a tool message's media inside its block). +MEDIA_ROLES: tuple[str, ...] = ("user", "tool") diff --git a/cosmos_framework/utils/generator/reasoner/model_path.py b/cosmos_framework/utils/generator/reasoner/model_path.py new file mode 100644 index 000000000..a4d11fe4f --- /dev/null +++ b/cosmos_framework/utils/generator/reasoner/model_path.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Resolve staged VLM training checkpoints without changing converter downloads.""" + +from pathlib import Path + +from cosmos_framework.utils.generator.reasoner.pretrained_models_downloader import maybe_download_hf_model_from_s3 + + +def resolve_hf_model_path( + model_name_or_path: str, + credentials: str, + bucket: str, + *, + include_model_weights: bool, + require_s3_exists: bool = False, +) -> str: + """Use a staged local directory or retain the normal remote download contract.""" + local_path = Path(model_name_or_path).expanduser() + if local_path.is_dir(): + return str(local_path.resolve()) + return maybe_download_hf_model_from_s3( + model_name_or_path, + credentials, + bucket, + include_model_weights=include_model_weights, + require_s3_exists=require_s3_exists, + ) diff --git a/cosmos_framework/utils/generator/reasoner/true_packing.py b/cosmos_framework/utils/generator/reasoner/true_packing.py index b9f9bb5dd..5cc0def1b 100644 --- a/cosmos_framework/utils/generator/reasoner/true_packing.py +++ b/cosmos_framework/utils/generator/reasoner/true_packing.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -"""Dense Qwen3-VL position helpers used only by true sequence packing.""" +"""Shared position helpers used only by true sequence packing.""" from types import SimpleNamespace from typing import Any @@ -39,7 +39,7 @@ def build_packed_position_ids( image_grid_thw: torch.LongTensor | None = None, # [N_image,3] video_grid_thw: torch.LongTensor | None = None, # [N_video,3] ) -> torch.Tensor: - """Build per-sample-reset M-RoPE ids for one true-packed dense Qwen3-VL row on CPU. + """Build per-sample-reset M-RoPE ids for one true-packed Qwen3-VL dense/MoE row on CPU. The unmodified dense Qwen3-VL position routine already resets each row independently and consumes concatenated image/video grids in row-major order. Reconstructing a small @@ -52,10 +52,9 @@ def build_packed_position_ids( Returns: Position ids with shape ``[3,1,S_pad]``. """ - if config.model_type != "qwen3_vl": + if config.model_type not in {"qwen3_vl", "qwen3_vl_moe"}: raise NotImplementedError( - "True-packing position ids are implemented for dense 'qwen3_vl' only; " - f"got model_type={config.model_type!r}." + f"True-packing position ids require Qwen3-VL dense or MoE; got model_type={config.model_type!r}." ) if input_ids.ndim != 2 or input_ids.shape[0] != 1: raise ValueError(f"True packing requires input_ids with shape [1,S_pad], got {tuple(input_ids.shape)}") @@ -100,11 +99,17 @@ def _grid_rows(grid: torch.Tensor | None, name: str) -> list[tuple[int, int, int temporal * (height // spatial_merge_size) * (width // spatial_merge_size) for temporal, height, width in _grid_rows(image_grid_thw, "image_grid_thw") ] + video_rows = _grid_rows(video_grid_thw, "video_grid_thw") video_runs = [ (height // spatial_merge_size) * (width // spatial_merge_size) - for temporal, height, width in _grid_rows(video_grid_thw, "video_grid_thw") + for temporal, height, width in video_rows for _ in range(temporal) ] + video_boundaries = {0} + frame_count = 0 + for temporal, _, _ in video_rows: + frame_count += temporal + video_boundaries.add(frame_count) image_index = video_index = 0 for sample_index, segment in enumerate(torch.split(real_input_ids, seq_lens)): tokens = [int(token) for token in segment.tolist()] @@ -138,6 +143,10 @@ def _grid_rows(grid: torch.Tensor | None, name: str) -> list[tuple[int, int, int raise ValueError( f"sample {sample_index} vision-start token is followed by unsupported token {modality_token}" ) + if video_index not in video_boundaries: + raise ValueError( + f"True-packing visual injection mismatch: video crosses logical sample boundary after sample {sample_index}" + ) if image_index != len(image_runs) or video_index != len(video_runs): raise ValueError( "True-packing visual injection mismatch: every image/video grid row must be consumed in " @@ -185,9 +194,10 @@ def prepare_true_packing_batch_on_cpu(data: dict[str, Any], config: Any) -> None """ if data.get("true_packing") is not True: raise ValueError("prepare_true_packing_batch_on_cpu requires a true-packed batch") - if config.model_type != "qwen3_vl": + supported_model_types = {"qwen3_vl", "qwen3_vl_moe"} + if config.model_type not in supported_model_types: raise NotImplementedError( - f"True packing is validated only for dense qwen3_vl; got model_type={config.model_type!r}" + f"True packing is validated only for {sorted(supported_model_types)}; got model_type={config.model_type!r}" ) input_ids = data.get("input_ids") @@ -211,11 +221,6 @@ def prepare_true_packing_batch_on_cpu(data: dict[str, Any], config: Any) -> None video_grid_thw = data.get("video_grid_thw") if video_grid_thw is not None and not isinstance(video_grid_thw, torch.Tensor): raise TypeError("video_grid_thw must be a tensor") - if video_grid_thw is not None and bool((video_grid_thw[:, 0] != 1).any()): - raise NotImplementedError( - "True-packed multi-frame video is not enabled until padded-vs-packed GPU logit and " - "gradient parity are certified; use one-frame video grids or padded batching." - ) data["position_ids"] = build_packed_position_ids( config, @@ -224,4 +229,7 @@ def prepare_true_packing_batch_on_cpu(data: dict[str, Any], config: Any) -> None image_grid_thw=data.get("image_grid_thw"), video_grid_thw=video_grid_thw, ) # [3,1,S_pad] + # Varlen attention still uses cu_seq_lens. CP slicing and MoE routing + # need a separate real-token mask, including unsupervised prompt tokens. + data["attention_mask"] = (torch.arange(input_ids.shape[1]) < sum(seq_lens)).unsqueeze(0) data[TRUE_PACKING_CPU_PREPARED_KEY] = True diff --git a/cosmos_framework/utils/generator/structured_conversation.py b/cosmos_framework/utils/generator/structured_conversation.py new file mode 100644 index 000000000..b897d112e --- /dev/null +++ b/cosmos_framework/utils/generator/structured_conversation.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + + +"""Restore stored agentic fields and preserve per-turn supervision for Qwen templates.""" + +from __future__ import annotations + +import copy +from typing import Any + +import numpy as np + +REASONING_PREFILL_KEY = "_reasoning_prefix_in_prompt" + + +def restore_agentic_messages( + messages: list[dict[str, Any]], +) -> tuple[list[dict[str, Any]], list[dict[str, Any]] | None]: + """Extract the tool menu and restore observation roles without changing storage.""" + restored = copy.deepcopy(messages) + tools = restored[0].pop("tools", None) if restored else None + for message in restored: + if message.get("role") == "user" and message.get("tool_call_id"): + message["role"] = "tool" + return restored, tools + + +def assistant_body_spans(tokens: list[int] | np.ndarray, tokenizer: Any) -> list[tuple[int, int]]: + """Find assistant bodies, retaining a BPE token shared with source whitespace.""" + values = np.asarray(tokens) + start_id = tokenizer.convert_tokens_to_ids("<|im_start|>") + end_id = tokenizer.convert_tokens_to_ids("<|im_end|>") + role = tokenizer.encode("assistant", add_special_tokens=False) + header = tokenizer.encode("assistant\n", add_special_tokens=False) + starts = np.flatnonzero(values == start_id) + ends = np.flatnonzero(values == end_id) + if len(starts) != len(ends): + raise ValueError("Unbalanced chat boundaries while applying source supervision") + spans = [] + previous_end = -1 + for start, end in zip(starts, ends): + if not previous_end < start < end: + raise ValueError("Nested or reordered chat boundaries") + previous_end = end + body = int(start + 1 + len(header)) + if values[start + 1 : body].tolist() != header: + if values[start + 1 : start + 1 + len(role)].tolist() != role: + continue + candidate = values[start + 1 : min(body + 4, end)].tolist() + decoded = tokenizer.decode(candidate, skip_special_tokens=False, clean_up_tokenization_spaces=False) + if not decoded.startswith("assistant\n"): + raise ValueError("Unexpected assistant role header") + whole = 0 + for size in range(1, len(candidate) + 1): + partial = tokenizer.decode( + candidate[:size], skip_special_tokens=False, clean_up_tokenization_spaces=False + ) + if not "assistant\n".startswith(partial): + break + whole = size + if whole == 0: + raise ValueError("No complete assistant role token") + body = int(start + 1 + whole) + spans.append((body, int(end + 1))) + return spans + + +def restrict_assistant_loss( + tokens: list[int] | np.ndarray, + mask: list[bool] | np.ndarray, + messages: list[dict[str, Any]], + tokenizer: Any, +) -> np.ndarray: + """Only remove loss: context assistant turns and inserted reasoning prefills. + + Fail on mismatched boundaries instead of allowing a context turn into loss. + This runs after modality token expansion, using the final token sequence. + """ + values = np.asarray(tokens) + result = np.asarray(mask, dtype=bool).copy() + if values.ndim != 1 or result.shape != values.shape: + raise ValueError("Expected matching one-dimensional token and loss arrays") + assistants = [message for message in messages if message.get("role") == "assistant"] + if not any(message.get("loss") is False or message.get(REASONING_PREFILL_KEY) for message in assistants): + return result + spans = assistant_body_spans(values, tokenizer) + if len(spans) != len(assistants): + raise ValueError("Assistant message count differs from rendered assistant frames") + for (start, end), message in zip(spans, assistants): + if message.get("loss") is False: + result[start:end] = False + elif message.get(REASONING_PREFILL_KEY): + prefix_text = "\n" + prefix = tokenizer.encode(prefix_text, add_special_tokens=False) + if values[start : start + len(prefix)].tolist() == prefix: + result[start : start + len(prefix)] = False + continue + # BPE can merge the inserted newline with source-leading whitespace. + # Mask only tokens wholly inside the inserted prefix. Keep a shared + # boundary token supervised so no source reasoning token is lost. + candidate = values[start : min(start + len(prefix) + 4, end)].tolist() + decoded = tokenizer.decode(candidate, skip_special_tokens=False, clean_up_tokenization_spaces=False) + if not decoded.startswith(prefix_text): + raise ValueError("Structured reasoning prefill differs from tokenizer output") + prefix_tokens = 0 + for size in range(1, len(candidate) + 1): + partial = tokenizer.decode( + candidate[:size], skip_special_tokens=False, clean_up_tokenization_spaces=False + ) + if not prefix_text.startswith(partial): + break + prefix_tokens = size + if prefix_tokens == 0: + raise ValueError("No complete token inside structured reasoning prefill") + result[start : start + prefix_tokens] = False + return result diff --git a/cosmos_framework/utils/generator/torchcodec_dependencies.py b/cosmos_framework/utils/generator/torchcodec_dependencies.py index 5776ae8f2..008166628 100644 --- a/cosmos_framework/utils/generator/torchcodec_dependencies.py +++ b/cosmos_framework/utils/generator/torchcodec_dependencies.py @@ -39,7 +39,11 @@ def preload_npp(cuda_version: str | None) -> ctypes.CDLL | None: for path in package.files or (): if path.name == library: # NPP's own $ORIGIN runpath resolves its sibling libnppc dependency. - return ctypes.CDLL(str(package.locate_file(path))) + try: + return ctypes.CDLL(str(package.locate_file(path))) + except OSError: + # Leave the failure to TorchCodec's full loader diagnostics. + return None return None diff --git a/cosmos_framework/utils/reasoner/dcp_checkpointer.py b/cosmos_framework/utils/reasoner/dcp_checkpointer.py index 8ab8fe791..15c0c29df 100644 --- a/cosmos_framework/utils/reasoner/dcp_checkpointer.py +++ b/cosmos_framework/utils/reasoner/dcp_checkpointer.py @@ -55,8 +55,8 @@ from cosmos_framework.checkpoint.base import AbstractCheckpointer, CheckpointLoadSource from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader, S3StorageWriter -from cosmos_framework.utils import callback, distributed, log, misc from cosmos_framework.utils.config import CheckpointConfig, JobConfig +from cosmos_framework.utils import callback, distributed, log, misc from cosmos_framework.utils.easy_io import easy_io from cosmos_framework.utils.reasoner.model_wrapper import ModelWrapper from cosmos_framework.utils.reasoner.optimizer import OptimizersContainer @@ -330,9 +330,18 @@ def load( log.info("scheduler.load_state_dict done", rank0_only=False) elif key == "trainer": log.info(f"- Loading the trainer {cur_key_ckpt_full_path}...", rank0_only=False) - # Here we skip loading the trainer, since 1) we only need the iteration which could be parsed from the name 2) we dont use grad_scaler - iteration = cur_key_ckpt_full_path.split("iter_")[-1].split("/")[0] - iteration = int(iteration) + if os.path.basename(os.path.normpath(source.path)).startswith("epoch_"): + # Epoch-named checkpoints carry no iteration in their name; save() records it. + trainer_state = easy_io.load( + os.path.join(cur_key_ckpt_full_path, f"rank_{dist.get_rank()}.pkl"), + file_format="pkl", + backend_key=source.backend_key, + ) + iteration = int(trainer_state["iteration"]) + else: + # Here we skip loading the trainer, since 1) we only need the iteration which could be parsed from the name 2) we dont use grad_scaler + iteration = cur_key_ckpt_full_path.split("iter_")[-1].split("/")[0] + iteration = int(iteration) elif key == "dataloader": rank = dist.get_rank() dataloader_pkl_path = os.path.join(cur_key_ckpt_full_path, f"rank_{rank}.pkl") @@ -429,7 +438,7 @@ def save_state_dict_worker( checkpoint_file: str, ) -> None: for k, (v, full_checkpoint_path) in to_save_dict.items(): - if k == "dataloader": + if k in ("dataloader", "trainer"): self._save_as_pkl(v, full_checkpoint_path) else: log.info(f"Saving {k} checkpoint to {full_checkpoint_path}") @@ -479,6 +488,9 @@ def save( } self.callbacks.on_save_checkpoint(model_parts, state_dict=to_save_dict) + if epoch is not None: + # load() reads the iteration from an iter_* name; an epoch-named checkpoint records it. + to_save_dict["trainer"] = {"iteration": iteration} for k in to_save_dict.keys(): output_dirname = os.path.join(self.save_dirname, checkpoint_file, k) @@ -486,10 +498,11 @@ def save( if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM: # Dataloader state is a plain dict, not tensors; save in main process and exclude from async staging. - dataloader_entry = to_save_dict.pop("dataloader", None) - if dataloader_entry is not None: - dataloader_state, dataloader_save_dir = dataloader_entry - self._save_as_pkl(dataloader_state, dataloader_save_dir) + for pkl_key in ("dataloader", "trainer"): + pkl_entry = to_save_dict.pop(pkl_key, None) + if pkl_entry is not None: + pkl_state, pkl_save_dir = pkl_entry + self._save_as_pkl(pkl_state, pkl_save_dir) self._async_with_pinned_memory(checkpoint_file, to_save_dict) else: self.save_state_dict_worker(to_save_dict, checkpoint_file) From 0aa4a8515258f1669e40cf253011092544596e5f Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Sun, 25 Oct 2026 08:44:53 +0530 Subject: [PATCH 2/2] test: drop the i4-only public-alias round-trip test from the release distillation_export_test.py imported cosmos3.common.public_model_config, which only exists in i4, so unittest failed to collect the module. The test checks restore under the i4 projects/ layout and does not apply here. Restore the file to main; i4 now keeps the test out of the release, so the next release produces this file unchanged. Co-Authored-By: Claude Opus 5.5 --- .../common/distillation_export_test.py | 48 ------------------- 1 file changed, 48 deletions(-) diff --git a/cosmos_framework/inference/common/distillation_export_test.py b/cosmos_framework/inference/common/distillation_export_test.py index 3bcd50360..697f4a281 100644 --- a/cosmos_framework/inference/common/distillation_export_test.py +++ b/cosmos_framework/inference/common/distillation_export_test.py @@ -5,11 +5,6 @@ import attrs import pytest -from cosmos3.common import public_model_config -from cosmos3.common.public_model_config import ( - build_public_model_config, - restore_model_config_from_public_model_config, -) from cosmos_framework.inference.common import distillation_export from cosmos_framework.inference.common.distillation_export import ( @@ -326,49 +321,6 @@ def test_sanitize_student_public_model_config_preserves_qwen_variant(model_size: ) -def test_sanitized_student_config_round_trips_through_public_aliases(monkeypatch: pytest.MonkeyPatch) -> None: - model_dict = { - "_target_": "cosmos_framework.model.generator.omni_mot_model.OmniMoTModel", - "config": { - "_type": "cosmos_framework.configs.base.defaults.model_config.OmniMoTModelConfig", - "vlm_config": { - "_type": "cosmos_framework.configs.base.defaults.reasoner.VLMConfig", - "model_instance": { - "_target_": "cosmos_framework.model.generator.mot.unified_mot.Qwen3VLTextForCausalLM", - "config": { - "_target_": "cosmos_framework.configs.base.defaults.reasoner.create_vlm_config", - "base_config": { - "_target_": ( - "cosmos_framework.model.generator.mot.unified_mot.Qwen3VLMoTConfig.from_json_file" - ), - "json_file": "/private/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json", - }, - }, - }, - }, - }, - } - - distillation_export.sanitize_student_public_model_config(model_dict) - exported_model_config = build_public_model_config(model_dict) - - exported_base_config = exported_model_config["config"]["vlm_config"]["model_instance"]["config"]["base_config"] - assert exported_model_config["_target"] == "omni_mot_model" - assert exported_base_config["json_file"] == ("cosmos3://vfm/models/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json") - - monkeypatch.setattr( - public_model_config, - "_module_exists", - lambda module: module == "projects.cosmos3.cosmos3", - ) - restored_model_config = restore_model_config_from_public_model_config(exported_model_config) - - restored_base_config = restored_model_config["config"]["vlm_config"]["model_instance"]["config"]["base_config"] - assert restored_base_config["json_file"] == ( - "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" - ) - - def test_sanitize_student_public_model_config_is_idempotent() -> None: model_dict = { "config": {