From 3a6ae4f19d025a474c008370ebbcfeb63b6594ab Mon Sep 17 00:00:00 2001 From: Liang Feng Date: Mon, 26 Oct 2026 05:53:28 +0000 Subject: [PATCH] Release 2026-10-26 from i4 09e71fb667f046f9f81f3935f13cd3b2143fd576 _source_commit: 09e71fb667f046f9f81f3935f13cd3b2143fd576-dirty _dest_commit: 4f6e539f98d1e7951e1c0907e61e2ff3ccdbdc92 --- .file_mapping.json | 6 +++--- .../configs/base/defaults/model_config.py | 6 +++--- .../generator/mot/context_parallel_test.py | 9 +++++---- .../model/generator/omni_mot_model.py | 20 +++++++++---------- 4 files changed, 20 insertions(+), 21 deletions(-) diff --git a/.file_mapping.json b/.file_mapping.json index 588ab449..7522487c 100644 --- a/.file_mapping.json +++ b/.file_mapping.json @@ -1,7 +1,7 @@ { - "_source_commit": "b3866eb30721bf68dce73380563235d39fa5a622-dirty", - "_dest_commit": "bf492fe2fda2108aae9b9561dc9fb9eec3fe2027", - "_generated_at": "2026-10-24T15:53:46Z", + "_source_commit": "09e71fb667f046f9f81f3935f13cd3b2143fd576-dirty", + "_dest_commit": "4f6e539f98d1e7951e1c0907e61e2ff3ccdbdc92", + "_generated_at": "2026-10-26T05:53:07Z", "files": { "imaginaire/__init__.py": "cosmos_framework/__init__.py", "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py", diff --git a/cosmos_framework/configs/base/defaults/model_config.py b/cosmos_framework/configs/base/defaults/model_config.py index a9fb5730..ff81ba5f 100644 --- a/cosmos_framework/configs/base/defaults/model_config.py +++ b/cosmos_framework/configs/base/defaults/model_config.py @@ -264,9 +264,9 @@ class OmniMoTModelConfig: # whose noised tokens share one timestep, take the unchunked path. None disables chunking. timestep_embedding_chunk_size: int | None = 128_000 - # Keep the raw CP data window on the host and release raw video CUDA inputs - # immediately after slot-0 tokenization. H100 multiview experiments opt in; - # the model default remains disabled. + # Keep the raw CP data window on the host and release raw vision CUDA inputs + # immediately after slot-0 tokenization. Memory-constrained CP experiments + # opt in; the model default remains disabled. release_cp_raw_batch_after_preprocessing: bool = False # torch.compile knobs (enabled, compiled_region, dynamic, ...). diff --git a/cosmos_framework/model/generator/mot/context_parallel_test.py b/cosmos_framework/model/generator/mot/context_parallel_test.py index c9e4cd2d..148fc730 100644 --- a/cosmos_framework/model/generator/mot/context_parallel_test.py +++ b/cosmos_framework/model/generator/mot/context_parallel_test.py @@ -141,12 +141,13 @@ def test_context_parallel_fetch_is_release_unaware( @pytest.mark.L0 @pytest.mark.CPU -def test_next_step_batch_routes_release_mode_batches(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize("media_key", ["video", "images"]) +def test_next_step_batch_routes_release_mode_batches(monkeypatch: pytest.MonkeyPatch, media_key: str) -> None: """StepBatch encapsulates host retention, CUDA transfer, and callback merging.""" trainer, model = _make_window_trainer(cp_size=2) model.config = SimpleNamespace(release_cp_raw_batch_after_preprocessing=True) - raw_video = torch.zeros(1, 3, 2, 4, 4, dtype=torch.uint8) - host_batch = {"video": raw_video, "dataset_name": "host"} + raw_vision = torch.zeros(1, 3, 2, 4, 4, dtype=torch.uint8) # [B,C,T,H,W] + host_batch = {media_key: raw_vision, "dataset_name": "host"} dataloader_iter = _CountingIterator([host_batch]) monkeypatch.setattr(dist, "get_backend", lambda group: dist.Backend.GLOO) monkeypatch.setattr(dist, "all_reduce", lambda tensor, op, group: None) @@ -174,7 +175,7 @@ def test_next_step_batch_routes_release_mode_batches(monkeypatch: pytest.MonkeyP assert not step1._is_fetch_slot assert step1.training is prepared cb = step1.callback_post_step - assert cb["video"] is raw_video + assert cb[media_key] is raw_vision assert cb["packed_tokens"] is prepared["packed_tokens"] assert cb["dataset_name"] == "prepared" assert dataloader_iter.fetch_count == 1 diff --git a/cosmos_framework/model/generator/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py index e46666a0..4bee26de 100644 --- a/cosmos_framework/model/generator/omni_mot_model.py +++ b/cosmos_framework/model/generator/omni_mot_model.py @@ -1409,29 +1409,27 @@ def _unpack_training_payload( vae_pixel_shapes = payload["vae_pixel_shapes"] return input_text_indexes, sequence_plans, gen_data_clean, memory_info, data_resolutions, vae_pixel_shapes - def _release_raw_video_after_tokenization( + def _release_raw_vision_after_tokenization( self, gen_data_clean: GenerationDataClean, data_batch: dict[str, torch.Tensor], ) -> None: - """Free raw video tensors from GPU after CP slot-0 tokenization. + """Free raw image or video tensors from GPU after CP slot-0 tokenization. - Only video (the largest modality) is dropped, and only for non-image - batches. The trainer retains the original host batch for callbacks, - which re-encode from its uint8 video. + The trainer retains the original host batch for callbacks, which + re-encode from its image or video tensors. ``num_vision_items_per_sample`` and ``image_size`` are kept: both describe the flattened vision-item layout and must stay mutually - consistent. :meth:`build_callback_batch` restores host video and - flattens it to match the retained counts so callbacks see a - self-consistent batch without recomputation. + consistent. :meth:`build_callback_batch` restores host vision and + flattens it to match the retained counts so callbacks see a self-consistent + batch without recomputation. """ if not self.config.release_cp_raw_batch_after_preprocessing: return - if gen_data_clean.is_image_batch: - return gen_data_clean.raw_state_vision = None data_batch.pop(self.input_video_key, None) + data_batch.pop(self.input_image_key, None) def _get_training_inputs( self, data_batch: dict[str, torch.Tensor], iteration: int @@ -1476,7 +1474,7 @@ def _get_training_inputs( if cp_window_slot == 0: local_training_data = self._prepare_training_data(data_batch, iteration) self._cp_local_training_payload = self._pack_training_payload(*local_training_data) - self._release_raw_video_after_tokenization(local_training_data[2], data_batch) + self._release_raw_vision_after_tokenization(local_training_data[2], data_batch) if self._cp_local_training_payload is None: raise RuntimeError("CP training payload cache is empty before its data window is complete.") payload = broadcast_context_parallel_object(