Spaces:
Running on L40S
Running on L40S
File size: 4,127 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | # 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
|