File size: 6,488 Bytes
9aa90e0 | 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 | # SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import os
from pathlib import Path
from typing import TYPE_CHECKING, NamedTuple
from loguru import logger
import ttnn
from ..layers.module import Module
from . import walltime
if TYPE_CHECKING:
from collections.abc import Callable, Sequence
CACHE_DICT_FILE = "cache_dict.json"
class MissingCacheError(Exception):
def __init__(self, path: str | Path) -> None:
self.path = Path(path)
def __str__(self) -> str:
return f"cache does not exist at '{self.path}'"
def config_id(parallel_config):
config_id = ""
for n, v in parallel_config._asdict().items():
if v is not None:
config_id += f"{''.join([w[0].upper() for w in n.split('_')])}{v.factor}_{v.mesh_axis}_"
return config_id
def cache_dir_is_set() -> bool:
return _cache_root() is not None
def load_model(
tt_model: Module,
*,
model_name: str,
subfolder: str,
parallel_config: NamedTuple,
mesh_shape: Sequence[int],
mesh_device: ttnn.MeshDevice,
dtype: str = "bf16",
is_fsdp: bool = False,
get_torch_state_dict: Callable[[], dict] | None = None,
create_cache: bool = True,
) -> None:
"""
Load model weights from cache or PyTorch state dict.
Attempts to load from cache first. If the cache does not exist, loads from PyTorch state dict
(if provided) and optionally creates the cache. Raises `MissingCacheError` if neither is
available. Finally, any module that needs to be offloaded is taken care of.
Args:
`tt_model`: TT model instance to load weights into.
`model_name`: Model name (e.g., "flux1-dev", "stable-diffusion-3.5").
`subfolder`: Subfolder within model cache directory (e.g., "transformer", "vae").
`parallel_config`: Parallelism configuration (tensor/sequence parallel).
`mesh_shape`: Device mesh shape.
`mesh_device`: Mesh device used to derive the multi-host ownership cache suffix.
`dtype`: Data type for cached weights (default: "bf16").
`is_fsdp`: Whether FSDP is used (default: False).
`get_torch_state_dict`: Optional callable returning PyTorch state dict. Enables lazy
evaluation - PyTorch model only loads if the cache does not exist. If `None`, cache
must exist or `MissingCacheError` is raised.
`create_cache`: Create cache after loading from PyTorch (default: True).
Raises:
`MissingCacheError`: Cache does not exist and `get_torch_state_dict` is `None`.
`RuntimeError`: `TT_DIT_CACHE_DIR` is not set and `get_torch_state_dict` is `None`.
"""
if tt_model.is_loaded():
return
cache_dir = model_cache_dir(
model_name=model_name,
subfolder=subfolder,
parallel_config=parallel_config,
mesh_shape=mesh_shape,
mesh_device=mesh_device,
dtype=dtype,
is_fsdp=is_fsdp,
required=get_torch_state_dict is None,
)
if cache_dir is None:
assert get_torch_state_dict is not None
logger.info(
"Loading transformer weights from PyTorch state dict. "
"To use caching, set the TT_DIT_CACHE_DIR environment variable."
)
with walltime.timed("weight_load", f"{model_name}/{subfolder}", cached=False):
tt_model.load_torch_state_dict(get_torch_state_dict())
ttnn.distributed_context_barrier()
return
if _cache_is_complete(cache_dir):
logger.info(f"loading cache at '{cache_dir}'.")
with walltime.timed("weight_load", f"{model_name}/{subfolder}", cached=True):
tt_model.load(cache_dir)
ttnn.distributed_context_barrier()
return
if get_torch_state_dict is None:
raise MissingCacheError(cache_dir)
logger.info("Cache does not exist. Loading PyTorch state dict.")
with walltime.timed("weight_load", f"{model_name}/{subfolder}", cached=False):
tt_model.load_torch_state_dict(get_torch_state_dict())
# If distributed, ensure that all processes have completed the check whether cache_dir exists,
# before any rank might proceed to create that dir to save.
ttnn.distributed_context_barrier()
if create_cache:
logger.info(f"Writing cache to '{cache_dir}'.")
tt_model.save(cache_dir)
_mark_cache_complete(cache_dir)
def model_cache_dir(
*,
model_name: str,
subfolder: str,
parallel_config: NamedTuple,
mesh_shape: Sequence[int],
mesh_device: ttnn.MeshDevice,
dtype: str = "bf16",
is_fsdp: bool = False,
required: bool = True,
) -> Path | None:
cache_dir = _cache_root()
if cache_dir is None:
if required:
msg = "Cache is required. Set the TT_DIT_CACHE_DIR environment variable."
raise RuntimeError(msg)
return None
parallel_key = config_id(parallel_config)
mesh_key = "x".join(str(x) for x in mesh_shape)
key = f"{parallel_key}mesh{mesh_key}_{dtype}"
if is_fsdp:
key += "_FSDP"
path = Path(cache_dir) / model_name / subfolder / key
ownership_suffix = _cache_ownership_suffix(mesh_device)
if ownership_suffix:
path = path / ownership_suffix
return path
def _cache_ownership_suffix(mesh_device: ttnn.MeshDevice) -> str:
"""Multi-host cache dir suffix keyed by local mesh-coordinate ownership.
Single-host / no distributed context: empty (same unsuffixed path as before).
Multi-host: ``host_coords_r{r0}-{r1}_c{c0}-{c1}`` for the local coord bounding box.
"""
if _distributed_world_size() <= 1:
return ""
view = mesh_device.get_view()
rows = []
cols = []
for coord in ttnn.MeshCoordinateRange(view.shape()):
if view.is_local(coord):
rows.append(int(coord[0]))
cols.append(int(coord[1]))
return f"host_coords_r{min(rows)}-{max(rows)}_c{min(cols)}-{max(cols)}"
def _cache_is_complete(cache_dir: str | Path) -> bool:
return (Path(cache_dir) / CACHE_DICT_FILE).is_file()
def _mark_cache_complete(cache_dir: str | Path) -> None:
(Path(cache_dir) / CACHE_DICT_FILE).touch()
def _distributed_world_size() -> int:
if not ttnn.distributed_context_is_initialized():
return 1
return int(ttnn.distributed_context_world_size())
def _cache_root() -> str | None:
return os.environ.get("TT_DIT_CACHE_DIR")
|