Spaces:
Build error
Build error
File size: 6,781 Bytes
b6af722 |
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 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import os
import threading
import torch
from torch._dynamo.eval_frame import OptimizedModule as torch_OptimizedModule
from cosmos_predict1.utils import callback, distributed, ema, log, misc
from cosmos_predict1.utils.checkpointer import Checkpointer
from cosmos_predict1.utils.config import CheckpointConfig, JobConfig
from cosmos_predict1.utils.model import Model
class TokenizerCheckpointer(Checkpointer):
"""The tokenizer checkpointer, extends the shared checkpointer.
Supports checkpoint saving/loading to local disk:
- network weights and training optimizer states.
- optionally, export a TorchScript version of the EMA model.
"""
def __init__(self, config_checkpoint: CheckpointConfig, config_job: JobConfig, callbacks: callback.CallBackGroup):
super().__init__(config_checkpoint, config_job, callbacks)
self.callbacks = callbacks
self.config_jit = config_checkpoint.jit
def save(
self,
model: Model,
optimizer: torch.optim.Optimizer,
scheduler: torch.optim.lr_scheduler.LRScheduler,
grad_scaler: torch.amp.GradScaler,
iteration: int = -1,
**ignore_kwargs,
) -> None:
"""Saves network weights, optimizer parameters, scheduler parameters to a checkpoint.
Args:
model (Model): The PyTorch model.
optimizer: The model optimizer.
scheduler: The optimization scheduler.
grad_scaler: The gradient scaler (for mixed precision training).
iteration: Current iteration number.
"""
self.callbacks.on_save_checkpoint_start(model, iteration)
model.eval()
checkpoint_file = f"iter_{iteration:09}.pt"
if distributed.get_rank() == 0:
state_dict = dict(
model=model.state_dict(),
optimizer=optimizer.state_dict(),
scheduler=scheduler.state_dict(),
grad_scaler=grad_scaler.state_dict(),
iteration=iteration,
)
state_dict = misc.to(state_dict, device="cpu")
self.callbacks.on_save_checkpoint(model, state_dict=state_dict)
# Wait for previous saver thread to end.
if self.save_thread:
self.save_thread.join()
# Run the checkpoint saver in a separate thread.
self.save_thread = threading.Thread(
target=self._save_worker_local,
daemon=False,
args=(state_dict, self._get_ema_jit(model), checkpoint_file, distributed.get_rank()),
)
self.save_thread.start()
# Note: Checkpoints are saved on a separate thread and this callback is not accurate.
# Please check logs from on_save_checkpoint_success() for better accuracy
self.callbacks.on_save_checkpoint_end(model=None, iteration=iteration)
@misc.timer("checkpoint saving (local)")
def _save_worker_local(
self,
state_dict: dict[str, torch.Tensor],
jit_models: dict[str, torch.ScriptModule],
checkpoint_file: str,
rank: int = 0,
) -> None:
"""Worker to save checkpoint to local disk, spawned with a child thread (runs in parallel with the training).
Args:
state_dict: The state dict of the model/optimizer/scheduler.
ema_jit: A dict of TorchScript EMA model, representing the encoder, decoder and full model.
checkpoint_file (str): The file name of the model checkpoint.
rank (int): GPU device (default: 0).
"""
checkpoint_path = os.path.join(self.checkpoint_dir_local, checkpoint_file)
os.makedirs(self.checkpoint_dir_local, exist_ok=True)
try:
torch.save(state_dict, checkpoint_path)
for key, jit_model in jit_models.items():
checkpoint_jit = checkpoint_path.replace(".pt", f"_{key}.jit")
torch.jit.save(jit_model, checkpoint_jit)
log.success(f"Saved checkpoint: {checkpoint_jit}")
if rank == 0:
self._write_latest_checkpoint_file(checkpoint_file)
log.success(f"Saved checkpoint (local): {checkpoint_path}")
iteration = int(checkpoint_file.replace("iter_", "").replace(".pt", ""))
self.callbacks.on_save_checkpoint_success(iteration=iteration)
except Exception as e: # noqa: BLE001
log.exception(f"Checkpoint failed to save (local): {e}")
def _get_ema_jit(self, model: Model) -> dict[str, torch.ScriptModule]:
"""Returns a TorchScript version of ema models compiled by PyTorch JIT."""
if not self.config_jit.enabled:
return dict()
input_shape = tuple(self.config_jit.input_shape)
example_input = torch.randn(input_shape)
dtype = getattr(torch, self.config_jit.dtype)
example_input = example_input.to(self.config_jit.device).to(dtype)
with ema.ema_scope(model, enabled=model.config.ema.enabled):
_model = model.network
if isinstance(_model, torch_OptimizedModule):
_model = _model._orig_mod
# Make sure jit model output consistenly during consecutive calls
# Check here: https://github.com/pytorch/pytorch/issues/74534
torch._C._jit_set_texpr_fuser_enabled(False)
ema_jit = torch.jit.trace(_model, example_input, strict=self.config_jit.strict)
encoder_jit = torch.jit.trace(_model.encoder_jit(), example_input, strict=self.config_jit.strict)
decoder_example = encoder_jit(example_input)
if isinstance(decoder_example, tuple):
decoder_example = decoder_example[0]
else:
assert isinstance(decoder_example, torch.Tensor), "decoder_example should be a tensor or tuple"
decoder_jit = torch.jit.trace(_model.decoder_jit(), decoder_example, strict=self.config_jit.strict)
return {"ema": ema_jit, "enc": encoder_jit, "dec": decoder_jit}
|