Spaces:
Running on L40S
Running on L40S
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: OpenMDW-1.1 | |
| """Training callback that defers AOT compilation of the VAE tokenizer. | |
| The actual compilation logic lives in | |
| :meth:`~projects.cosmos3.vfm.tokenizers.wan2pt2_vae_4x16x16.Wan2pt2VAEInterface.compile_encode`. | |
| This module provides a :class:`CompileTokenizer` callback that invokes it | |
| at the right point during training (after ``compile_after_iterations`` | |
| steps, to avoid NCCL timeouts during CUDA/cuDNN warm-up). | |
| Typical config usage | |
| -------------------- | |
| .. code-block:: python | |
| CompileTokenizer( | |
| enabled=True, | |
| compile_after_iterations=3, | |
| warmup_resolutions=["256", "480", "720"], | |
| ) | |
| """ | |
| from collections.abc import Sequence | |
| import torch | |
| from cosmos_framework.utils import log | |
| from cosmos_framework.utils.callback import Callback | |
| from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel | |
| class CompileTokenizer(Callback): | |
| """Training callback that defers AOT compilation of the VAE tokenizer. | |
| Hooks into ``on_training_step_start``. On the | |
| ``compile_after_iterations``-th step it calls | |
| ``Wan2pt2VAEInterface.compile_encode`` to compile and load all chunk | |
| variants. Every subsequent step is a no-op. | |
| """ | |
| def __init__( | |
| self, | |
| enabled: bool = False, | |
| compile_after_iterations: int = 3, | |
| warmup_resolutions: Sequence[str] | None = None, | |
| ): | |
| """ | |
| Args: | |
| enabled: Master switch. When ``False`` the callback is a | |
| complete no-op and no compilation occurs. | |
| compile_after_iterations: How many training steps to skip | |
| before triggering compilation. The default (3) lets CUDA | |
| context setup and Transformer compilation finish first. | |
| warmup_resolutions: Resolution keys (e.g. ``["256", "480", "720"]``) | |
| to AOT-compile. Should include every resolution used in | |
| training. Must be a non-empty list when *enabled* is ``True``. | |
| """ | |
| super().__init__() | |
| self.enabled: bool = enabled | |
| self.compile_after_iterations: int = compile_after_iterations | |
| self.skip_counter: int = 0 | |
| self.warmup_resolutions: Sequence[str] | None = warmup_resolutions | |
| if self.enabled: | |
| if self.warmup_resolutions is None: | |
| raise ValueError("warmup_resolutions must be provided when enabled, got None") | |
| if len(self.warmup_resolutions) == 0: | |
| raise ValueError("warmup_resolutions must be a non-empty list when enabled, got an empty list") | |
| def on_training_step_start( | |
| self, model: OmniMoTModel, data_batch: dict[str, torch.Tensor], iteration: int = 0 | |
| ) -> None: | |
| """Called at the start of every training step. | |
| On the ``compile_after_iterations``-th call, triggers AOT compilation | |
| via ``tokenizer.compile_encode``. | |
| Args: | |
| model: The OmniMoTModel whose ``tokenizer_vision_gen`` will be compiled. | |
| data_batch: Current training batch (unused, required by Callback API). | |
| iteration: Current training iteration (unused; we track our own counter | |
| via ``skip_counter`` because this callback may be registered after | |
| iteration 0). | |
| """ | |
| if not self.enabled: | |
| return | |
| tokenizer = model.tokenizer_vision_gen | |
| if isinstance(tokenizer, torch.jit.ScriptModule): | |
| log.critical( | |
| f"The Tokenizer model {type(tokenizer)} is a JIT model, " | |
| "which is not compilable. The Tokenizer will not be compiled.", | |
| rank0_only=False, | |
| ) | |
| self.enabled = False | |
| return | |
| if self.skip_counter == self.compile_after_iterations: | |
| if self.warmup_resolutions is not None: | |
| tokenizer.compile_encode( | |
| self.warmup_resolutions, | |
| output_dir=self.config.job.path_local, | |
| ) | |
| self.skip_counter += 1 | |