File size: 13,521 Bytes
b025706 | 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 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 | # SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
"""
LazyWeight: Lazy weight loading with explicit TTNN conversion parameters.
This module has NO torch dependency - it accepts any tensor-like object that
ttnn.from_torch() can handle (duck typing with string type hints).
Design principles:
- Explicit parameters over hidden closures (IDE-friendly)
- Duck typing for source tensors (no torch import)
- Deterministic cache fingerprinting from configuration
- Lazy conversion: defer ttnn.from_torch until first access
- Cache-first: skip torch entirely if TTNN cache exists
"""
import hashlib
import pathlib
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Optional
from loguru import logger
import ttnn
from models.common.tensor_utils import get_padded_hidden_dim, pad_to_shape, parse_shard_dims_from_mesh_mapper_config
# todo)) maybe useful to support a mechanism that can be used to disable the cache for every LazyWeight instance
# todo)) needs unit tests to provide coverage
@dataclass
class LazyWeight:
"""
Lazy weight loading with explicit TTNN conversion parameters.
The source can be:
- A tensor directly (duck-typed, typically torch.Tensor)
All fields except source are optional at construction time and can be
provided later. get_weight() will check
Example usage:
# Fully specified at construction
weight = LazyWeight(
source=torch_tensor,
device=mesh_device,
dtype=ttnn.bfloat16,
)
ttnn_tensor = weight.get_weight()
# Minimal construction, provide defaults at materialize time
weight = LazyWeight(source=torch_tensor)
ttnn_tensor = weight.get_weight(device=mesh_device, dtype=ttnn.bfloat16)
# Partial specification - override defaults selectively
weight = LazyWeight(source=torch_tensor, dtype=ttnn.bfloat4_b)
ttnn_tensor = weight.get_weight(device=mesh_device, dtype=ttnn.bfloat16) # uses bfloat4_b
# Using convenience class methods
weight = LazyWeight.sharded_1d(source=tensor, device=mesh_device, dtype=ttnn.bfloat16, dim=3)
weight = LazyWeight.replicated(source=tensor, device=mesh_device, dtype=ttnn.bfloat16)
"""
# Source: duck-typed tensor
# String annotation documents intent without importing torch
# todo)) use something like LazyStateDict for the source
source: "torch.Tensor"
# All other fields are optional at construction time.
cache_dir_weight_name: Optional[tuple[Path, str]] = None # do not cache if None
pad_value: Optional[float] = 0.0
dtype: Optional[ttnn.DataType] = ttnn.bfloat16
# Lazy fields that can be materialize at get_weight time;
# Still named as public fields to allow users to override at construction time (be proactive if you want).
device: Optional[ttnn.MeshDevice] = None
mesh_mapper_config: Optional[ttnn.MeshMapperConfig] = None
memory_config: Optional[ttnn.MemoryConfig] = None
layout: Optional[ttnn.Layout] = None
# Private fields
_value: Optional[ttnn.Tensor] = field(default=None, repr=False)
def __post_init__(self):
assert self.source.shape is not None and len(self.source.shape) > 0, "source must have a shape"
def get_device_weight(self) -> ttnn.Tensor:
"""
Get the TTNN tensor, converting from source if needed.
On first call:
1. Check cache → load if exists
2. Otherwise: resolve source → convert → cache
Subsequent calls return the cached tensor (defaults ignored after first call).
"""
# todo)) is there a better way to enforce immutability of all fields after construction? and self._value after first call?
if self._value is not None:
return self._value
cache_dir, weight_name = None, None
if self.cache_dir_weight_name is not None:
cache_dir, weight_name = self.cache_dir_weight_name
# Validate required fields
if self.device is None:
raise ValueError("device must be provided (either on LazyWeight or in defaults)")
if self.layout is None:
raise ValueError("layout must be provided (either on LazyWeight or in defaults)")
if self.memory_config is None:
raise ValueError("memory_config must be provided (either on LazyWeight or in defaults)")
# Try loading from cache first (skip torch entirely if possible)
cache_file_name = self._get_cache_fill_path(
cache_dir=cache_dir,
weight_name=weight_name,
)
if cache_file_name and cache_file_name.exists():
logger.info(f"\033[32m[cache hit]\033[0m Loading tensor from cache: {cache_file_name}")
self._value = ttnn.load_tensor(str(cache_file_name), device=self.device)
return self._value
# Resolve source
tensor = self.source
# Get mesh mapper (created from config)
if self.mesh_mapper_config is not None:
mesh_mapper = ttnn.create_mesh_mapper(self.device, self.mesh_mapper_config)
# Auto-pad tensor for tile alignment (only needed for TILE_LAYOUT sharding)
if self.layout == ttnn.TILE_LAYOUT:
tensor = _auto_pad_for_sharded_tiles(tensor, self.padded_shape, pad_value=self.pad_value)
else:
# None config means replicate
mesh_mapper = ttnn.replicate_tensor_to_mesh_mapper(self.device)
is_replicated = self.mesh_mapper_config is None
# Convert to TTNN and cache
self._value = _from_torch_and_dump(
tensor=tensor,
device=self.device,
dtype=self.dtype,
layout=self.layout,
memory_config=self.memory_config,
mesh_mapper=mesh_mapper,
is_replicated=is_replicated,
pad_value=self.pad_value,
cache_file_name=cache_file_name,
)
return self._value
# todo)) refactor into a base class for Lazy and Transparent Pattern (LaTr)
def is_resolved(self) -> bool:
"""Check if all required fields for weight materialization are set.
Excludes: _value (private), cache_dir_weight_name (optional caching)
"""
required_fields = ("source", "pad_value", "dtype", "device", "mesh_mapper_config", "memory_config", "layout")
return all(getattr(self, field) is not None for field in required_fields)
def _get_cache_fill_path(
self,
cache_dir: Optional[Path],
weight_name: Optional[str],
) -> Optional[Path]:
"""Generate the cache file path based on configuration fingerprint."""
if cache_dir is None or weight_name is None:
return None
fingerprint = self._get_fingerprint()
return cache_dir / f"{weight_name}_{fingerprint}.tensorbin"
def _get_fingerprint(
self,
) -> str:
"""
Generate a unique fingerprint from the conversion configuration.
This ensures cache invalidation when config changes.
"""
parts = []
# source
# todo)) add better fingerprinting for the source tensor to enable cache invalidation when the source tensor changes; for now, use shape.
# maybe a sampled hash (on, say, first N bytes) would be a good solution --> needs testing! --> test_rope_1d.py is a good test case with differring rope_theta values
parts.append(f"srcshape_{'_'.join(str(dim) for dim in self.source.shape)}")
# dtype
parts.append(f"dtype_{self.dtype.name}")
# layout
parts.append(f"layout_{self.layout.name}")
# memory_config (use hash if non-default)
if self.memory_config is not None and self.memory_config != ttnn.DRAM_MEMORY_CONFIG:
parts.append(f"memcfg_{hash(self.memory_config)}")
# mesh_mapper_config fingerprint
parts.append(self._get_mesh_mapper_fingerprint())
# device fingerprint
device_id = self.device.id() if hasattr(self.device, "id") else "single"
parts.append(f"device_{device_id}")
# pad_value (only if non-default)
if self.pad_value != 0.0:
parts.append(f"pad_{self.pad_value}")
return "_".join(parts)
def _get_mesh_mapper_fingerprint(self) -> str:
"""
Generate a fingerprint for the mesh_mapper_config.
Uses the config's __repr__ for deterministic hashing since
MeshMapperConfig exposes a proper string representation via pybind.
"""
if self.mesh_mapper_config is None:
return "replicated"
# Use repr of config for deterministic fingerprint
config_repr = repr(self.mesh_mapper_config)
# Hash to keep fingerprint short but unique
config_hash = hashlib.md5(config_repr.encode()).hexdigest()[:12]
return f"mapper_{config_hash}"
@property
def padded_shape(self) -> tuple[int, ...]:
"""
Returns the shape after padding for tile alignment.
Requires device and mesh_mapper_config to be set.
Note: source.shape remains the canonical unpadded shape.
Tile-alignment padding is only applied for TILE_LAYOUT. ROW_MAJOR_LAYOUT
does not require tile-aligned shard sizes.
"""
assert self.is_resolved(), "LazyWeight must be resolved to compute padded_shape"
if self.device is None:
raise ValueError("device must be set to compute padded_shape")
shape = list(self.source.shape)
if self.mesh_mapper_config is None:
return tuple(shape) # replicated, no padding
num_devices = self.device.get_num_devices()
if num_devices == 1:
return tuple(shape)
# ROW_MAJOR_LAYOUT doesn't require tile-aligned shard sizes
if self.layout == ttnn.ROW_MAJOR_LAYOUT:
return tuple(shape)
shard_dims = parse_shard_dims_from_mesh_mapper_config(self.mesh_mapper_config)
for shard_dim in shard_dims:
if shard_dim < 0:
shard_dim = len(shape) + shard_dim
shape[shard_dim] = get_padded_hidden_dim(shape[shard_dim], num_devices)
return tuple(shape)
def _auto_pad_for_sharded_tiles(
tensor: "torch.Tensor",
padded_shape: tuple[int, ...],
pad_value: float = 0.0,
) -> "torch.Tensor":
"""
Auto-pad tensor for tile-aligned sharding (TILE_LAYOUT only).
ttnn.from_torch requires physical shard shapes to be tile-aligned when using
TILE_LAYOUT. This function pads the global tensor to the pre-computed padded_shape.
Args:
tensor: Source torch tensor
padded_shape: Target shape after padding (from LazyWeight.padded_shape)
pad_value: Value to use for padding (default 0.0)
Returns:
Padded tensor if padding was needed, otherwise original tensor
"""
if tensor.shape != padded_shape:
logger.debug(f"Auto-padding from {tuple(tensor.shape)} to {padded_shape} for tile alignment")
return pad_to_shape(tensor, padded_shape, pad_value=pad_value)
return tensor
def _from_torch_and_dump(
tensor: "torch.Tensor",
device: Optional[ttnn.MeshDevice],
dtype: Optional[ttnn.DataType],
layout: Optional[ttnn.Layout],
memory_config: Optional[ttnn.MemoryConfig],
mesh_mapper: Optional[ttnn.CppTensorToMesh],
is_replicated: bool,
pad_value: float,
cache_file_name: Optional[str],
):
"""
Convert a torch tensor to TTNN format and optionally cache it.
Args:
tensor: Source torch tensor
device: Target MeshDevice
dtype: Target data type
layout: Target layout
memory_config: Target memory config
mesh_mapper: The mesh mapper to use for distribution
is_replicated: If True, cache the unsharded tensor for portability
cache_file_name: Path to cache file, or None to skip caching
"""
local_mesh_mapper = None if is_replicated else mesh_mapper
if cache_file_name is None:
return ttnn.from_torch(
tensor,
dtype=dtype,
layout=layout,
mesh_mapper=local_mesh_mapper,
memory_config=memory_config,
device=device,
pad_value=pad_value,
)
tensor = ttnn.from_torch(
tensor,
dtype=dtype,
layout=layout,
mesh_mapper=local_mesh_mapper,
memory_config=memory_config,
# For fully replicated tensors, cache unsharded tensor so it can be loaded on any device.
device=None,
pad_value=pad_value,
)
assert tensor.storage_type() == ttnn.StorageType.HOST, "tensor should be on host"
assert cache_file_name is not None, "cache_file_name must be provided for caching"
logger.debug(f"\033[33m[cache miss]\033[0m Generating cache for {cache_file_name}")
pathlib.Path(cache_file_name).parent.mkdir(parents=True, exist_ok=True)
ttnn._ttnn.tensor.dump_tensor_flatbuffer(str(cache_file_name), tensor)
if device is not None:
tensor = tensor.to(device, memory_config)
return tensor
def resolve_lazy_weight(weight: LazyWeight, **kwargs) -> LazyWeight:
"""Resolve the None fields of `weight` with the given kwargs; do not override non-None fields"""
to_set = {k: v for k, v in kwargs.items() if getattr(weight, k, None) is None}
return replace(weight, **to_set)
|