Spaces:
Running on L40S
Running on L40S
File size: 8,568 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 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 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
from __future__ import annotations
import os
from dataclasses import dataclass
from typing import Any
import torch
from cosmos_framework.model._base import ImaginaireModel
from cosmos_framework.utils import log
from cosmos_framework.utils.callback import Callback
@dataclass
class NoReplaceShardlistState:
epoch: int = 0
index: int = 0
class DataLoaderStateCallback(Callback):
checkpoint_component: str = "dataloader"
def __init__(
self,
distributor_type: str | None = None,
name: str = "",
) -> None:
super().__init__()
self.distributor_type = distributor_type
self.name = name
self.config: Any = None
self.state: dict[int, NoReplaceShardlistState] = {}
self.verbose = True
def _update_state_from_batch(self, data_batch: dict[str, torch.Tensor]) -> None:
if "sample_worker_id" not in data_batch:
return # batch has no position metadata (shuffle=False or iterable data_source)
worker_ids = data_batch["sample_worker_id"].tolist() # [B]
epochs = data_batch["sample_epoch"].tolist() # [B]
indices = data_batch["sample_index"].tolist() # [B]
for worker_id, epoch, index in zip(worker_ids, epochs, indices, strict=True):
if worker_id not in self.state:
self.state[worker_id] = NoReplaceShardlistState(epoch=epoch, index=index)
elif self.state[worker_id].epoch < epoch or (
self.state[worker_id].index < index and self.state[worker_id].epoch == epoch
):
self.state[worker_id] = NoReplaceShardlistState(epoch=epoch, index=index)
_ACTIVE_DISTRIBUTOR_TYPES = ("no_replace", "data_packer")
def on_training_step_batch_end(
self,
model: ImaginaireModel,
data_batch: dict[str, torch.Tensor],
output_batch: dict[str, torch.Tensor],
loss: torch.Tensor,
iteration: int = 0,
) -> None:
if self.distributor_type in self._ACTIVE_DISTRIBUTOR_TYPES:
self._update_state_from_batch(data_batch)
def on_training_step_end(
self,
model: ImaginaireModel,
data_batch: dict[str, torch.Tensor],
output_batch: dict[str, torch.Tensor],
loss: torch.Tensor,
iteration: int = 0,
) -> None:
if self.distributor_type in self._ACTIVE_DISTRIBUTOR_TYPES:
if self.verbose:
if iteration % self.config.trainer.logging_iter == 0:
msg = "\n"
for wid, state in self.state.items():
msg += f"worker {wid}: epoch={state.epoch}, index={state.index}\n"
log.info(msg)
def has_checkpoint_state(self) -> bool:
return self.distributor_type in self._ACTIVE_DISTRIBUTOR_TYPES
def state_dict(self) -> dict[int, dict[str, int]]:
if self.distributor_type not in self._ACTIVE_DISTRIBUTOR_TYPES:
return {}
state_dict: dict[int, dict[str, int]] = {}
for worker_id, per_worker_state in self.state.items():
state_dict[worker_id] = {"epoch": per_worker_state.epoch, "index": per_worker_state.index}
log.info(
f"Saved dataloader state for worker {worker_id}: "
f"epoch={per_worker_state.epoch}, index={per_worker_state.index}"
)
return state_dict
def load_state_dict(self, state_dict: dict[int, dict[str, int]]) -> None:
if self.distributor_type not in self._ACTIVE_DISTRIBUTOR_TYPES:
return
if not state_dict:
log.info("No dataloader state found in checkpoint")
return
self.state = {}
# Build env var prefix. For data_packer, namespacing avoids conflicts
# when multiple DataPackerDataLoader instances share the same process
# (e.g. inside JointDataPackerDataLoader). name="" → original format.
_dp_pfx = f"DP_STATE_{self.name}_" if self.name else "DP_STATE_"
for worker_id, per_worker_state in state_dict.items():
epoch = per_worker_state["epoch"]
index = per_worker_state["index"]
self.state[worker_id] = NoReplaceShardlistState(epoch=epoch, index=index)
if self.distributor_type == "data_packer":
os.environ[f"{_dp_pfx}WORKER_{worker_id}_EPOCH"] = str(epoch)
os.environ[f"{_dp_pfx}WORKER_{worker_id}_INDEX"] = str(index)
log.info(f"Loaded data_packer dataloader state for worker {worker_id}: epoch={epoch}, index={index}")
else:
os.environ[f"NSL_STATE_WORKER_{worker_id}_EPOCH"] = str(epoch)
os.environ[f"NSL_STATE_WORKER_{worker_id}_INDEX"] = str(index)
log.info(f"Loaded no_replace dataloader state for worker {worker_id}: epoch={epoch}, index={index}")
class JointDataLoaderStateCallback(Callback):
"""Checkpoint/resume state for ``JointDataPackerDataLoader``.
Manages two levels of state in a single DCP checkpoint entry
(``checkpoint_component = "dataloader"``):
1. **Outer** ``global_id`` — the number of batches the outer loader has
yielded. Restored via ``outer_loader.set_start_iteration(global_id)``
so the deterministic dataset-selection sequence resumes from the correct
step.
2. **Inner** per-dataset, per-worker ``(epoch, index)`` — one
``DataLoaderStateCallback`` per inner loader, keyed by the dataset name.
Each inner callback sets namespaced env vars on ``load_state_dict`` so
workers fast-forward to the saved sample position.
Usage in experiment configs::
joint_loader = JointDataPackerDataLoader(dataloaders={...}, seed=42)
exp["dataloader_train"] = joint_loader
exp["trainer"]["callbacks"]["dataloader_state"] = JointDataLoaderStateCallback(
outer_loader=joint_loader,
distributor_type="data_packer",
)
The ``checkpoint_component = "dataloader"`` class attribute ensures the DCP
checkpointer's ``_DataloaderWrapper`` discovers exactly this callback (it
picks the first matching callback). Do **not** also register standalone
``DataLoaderStateCallback`` instances for the inner loaders — this class
already handles them all.
"""
checkpoint_component: str = "dataloader"
def __init__(
self,
outer_loader: Any,
distributor_type: str = "data_packer",
) -> None:
super().__init__()
self._outer = outer_loader
self._inner: dict[str, DataLoaderStateCallback] = {
name: DataLoaderStateCallback(distributor_type=distributor_type, name=name)
for name in outer_loader._names
}
self.config: Any = None
def _update_state_from_batch(self, batch: dict) -> None:
name = batch.get("dataset_name")
if name in self._inner:
self._inner[name]._update_state_from_batch(batch)
def on_training_step_batch_end(
self,
model: Any,
data_batch: dict,
output_batch: dict,
loss: Any,
iteration: int = 0,
) -> None:
self._update_state_from_batch(data_batch)
def on_training_step_end(
self,
model: Any,
data_batch: dict,
output_batch: dict,
loss: Any,
iteration: int = 0,
) -> None:
if self.config and iteration % self.config.trainer.logging_iter == 0:
msg = f"\nJointDataPackerDataLoader global_id={self._outer._global_id}\n"
for name, cb in self._inner.items():
for wid, state in cb.state.items():
msg += f" [{name}] worker {wid}: epoch={state.epoch}, index={state.index}\n"
log.info(msg)
def has_checkpoint_state(self) -> bool:
return True
def state_dict(self) -> dict:
return {
"global_id": self._outer._global_id,
**{name: cb.state_dict() for name, cb in self._inner.items()},
}
def load_state_dict(self, state: dict) -> None:
global_id = state.get("global_id", 0)
self._outer.set_start_iteration(global_id)
log.info(f"JointDataLoaderStateCallback: resumed outer global_id={global_id}")
for name, cb in self._inner.items():
if name in state:
cb.load_state_dict(state[name])
|