Coverage for transformer_lens/utilities/multi_gpu.py: 87%
142 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Multi-GPU utilities.
3Utilities for managing multiple GPU devices and distributing model layers across them.
4"""
6from __future__ import annotations
8from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
10import torch
11from torch import nn
13if TYPE_CHECKING:
14 from transformer_lens.config.transformer_bridge_config import (
15 TransformerBridgeConfig as ConfigType,
16 )
17 from transformer_lens.config.transformer_lens_config import (
18 TransformerLensConfig as BaseConfigType,
19 )
20else:
21 ConfigType = Any
22 BaseConfigType = Any
24AvailableDeviceMemory = list[tuple[int, int]]
25"""
26This type is passed around between different CUDA memory operations.
27The first entry of each tuple will be the device index.
28The second entry will be how much memory is currently available.
29"""
31MaxMemory = Dict[Union[str, int], Union[str, int]]
34def calculate_available_device_cuda_memory(i: int) -> int:
35 """Calculates how much memory is available at this moment for the device at the indicated index
37 Args:
38 i (int): The index we are looking at
40 Returns:
41 int: How memory is available
42 """
43 total = torch.cuda.get_device_properties(i).total_memory
44 allocated = torch.cuda.memory_allocated(i)
45 return total - allocated
48def _get_available_cuda_memory_for_device(i: int) -> int:
49 """Return a concrete byte budget accepted by Accelerate's ``max_memory``."""
50 try:
51 free_memory, _ = torch.cuda.mem_get_info(i)
52 except (AttributeError, RuntimeError):
53 return calculate_available_device_cuda_memory(i)
54 return int(free_memory)
57def determine_available_memory_for_available_devices(max_devices: int) -> AvailableDeviceMemory:
58 """Gets all available CUDA devices with their current memory calculated
60 Returns:
61 AvailableDeviceMemory: The list of all available devices with memory precalculated
62 """
63 devices = []
64 for i in range(max_devices):
65 devices.append((i, calculate_available_device_cuda_memory(i)))
67 return devices
70def sort_devices_based_on_available_memory(devices: AvailableDeviceMemory) -> AvailableDeviceMemory:
71 """Sorts all available devices with devices with the most available memory returned first
73 Args:
74 devices (AvailableDeviceMemory): All available devices with memory calculated
76 Returns:
77 AvailableDeviceMemory: The same list of passed through devices sorted with devices with most
78 available memory first
79 """
80 return sorted(devices, key=lambda x: x[1], reverse=True)
83def get_best_available_cuda_device(max_devices: Optional[int] = None) -> torch.device:
84 """Gets whichever cuda device has the most available amount of memory for use
86 Raises:
87 EnvironmentError: If there are no available devices, this will error out
89 Returns:
90 torch.device: The specific device that should be used
91 """
92 max_devices = max_devices if max_devices is not None else torch.cuda.device_count()
93 devices = determine_available_memory_for_available_devices(max_devices)
95 if len(devices) <= 0:
96 raise EnvironmentError(
97 "TransformerLens has been configured to use CUDA, but no available devices are present"
98 )
100 sorted_devices = sort_devices_based_on_available_memory(devices=devices)
102 return torch.device("cuda", sorted_devices[0][0])
105def get_best_available_device(
106 cfg: ConfigType,
107) -> torch.device:
108 """Gets the best available device to be used based on the passed in arguments
110 Args:
111 cfg: The bridge config object containing device configuration
113 Returns:
114 torch.device: The best available device
115 """
116 assert cfg.device is not None
117 device = torch.device(cfg.device)
119 if device.type == "cuda" and cfg.n_devices > 1:
120 return get_best_available_cuda_device(cfg.n_devices)
121 else:
122 return device
125def get_device_for_block_index(
126 index: int,
127 cfg: "BaseConfigType",
128 device: Optional[Union[torch.device, str]] = None,
129):
130 """
131 Determine the device for a given layer index based on the model configuration.
133 This function assists in distributing model layers across multiple devices. The distribution
134 is based on the configuration's number of layers (cfg.n_layers) and devices (cfg.n_devices).
137 Args:
138 index (int): Model layer index.
139 cfg: Model and device configuration.
140 device (Optional[Union[torch.device, str]], optional): Initial device used for determining the target device.
141 If not provided, the function uses the device specified in the configuration (cfg.device).
143 Returns:
144 torch.device: The device for the specified layer index.
146 Deprecated:
147 This function did not take into account a few factors for multi-GPU support. You should now
148 use get_best_available_device in order to properly run models on multiple devices.
149 This will be removed in 3.0
150 """
151 assert cfg.device is not None
152 if device is None:
153 device = cfg.device
154 device = torch.device(device)
155 if device.type == "cpu":
156 return device
157 # Multiplying first guarantees the result is in [0, n_devices - 1] and avoids
158 # the divide-by-zero when n_layers < n_devices. The naive form
159 # `index // (n_layers // n_devices)` floors the divisor and overshoots when
160 # n_layers is not a multiple of n_devices (e.g. 62 layers / 8 devices → 8).
161 n_devices = getattr(cfg, "n_devices", 1)
162 device_index = (device.index or 0) + (index * n_devices) // cfg.n_layers
163 return torch.device(device.type, device_index)
166def resolve_device_map(
167 n_devices: Optional[int],
168 device_map: Optional[Union[str, Dict[str, Union[str, int]]]],
169 device: Optional[Union[str, torch.device]],
170 max_memory: Optional[MaxMemory] = None,
171) -> Tuple[Optional[Union[str, Dict[str, Union[str, int]]]], Optional[MaxMemory]]:
172 """Resolve ``n_devices`` / ``device_map`` / ``device`` into HF ``from_pretrained`` kwargs.
174 Returns ``(device_map, max_memory)`` tuple ready to pass into ``model_kwargs``.
176 Semantics:
177 - Explicit ``device_map`` wins and is passed through unchanged (user-provided
178 ``max_memory`` is passed through too). CPU targets are supported; disk / meta
179 offload targets are still rejected because Bridge component wrappers can bypass
180 Accelerate's offload hooks during forward passes.
181 - ``n_devices=None`` or ``1``: returns ``(None, None)`` — single-device path.
182 - ``n_devices > 1``: returns ``("balanced", {0: bytes, ..., n-1: bytes})``.
183 ``"balanced"`` is accelerate's string directive for balanced layer dispatch;
184 concrete byte budgets cap visibility to exactly ``n_devices`` GPUs.
185 """
186 if device_map is not None and device is not None:
187 raise ValueError("device and device_map are mutually exclusive — pass one.")
188 if device_map is not None:
189 _validate_device_map_values(device_map)
190 return device_map, max_memory
191 if n_devices is None or n_devices <= 1:
192 return None, max_memory
193 if not torch.cuda.is_available():
194 raise ValueError(f"n_devices={n_devices} requires CUDA, which is not available.")
195 if torch.cuda.device_count() < n_devices: 195 ↛ 196line 195 didn't jump to line 196 because the condition on line 195 was never true
196 raise ValueError(
197 f"n_devices={n_devices} but only {torch.cuda.device_count()} CUDA devices present."
198 )
199 resolved_max_memory: MaxMemory = (
200 dict(max_memory)
201 if max_memory is not None
202 else {i: _get_available_cuda_memory_for_device(i) for i in range(n_devices)}
203 )
204 return "balanced", resolved_max_memory
207def _validate_device_map_values(
208 device_map: Union[str, Dict[str, Union[str, int]]],
209) -> None:
210 """Reject mixed CPU/disk + GPU targets in a user-supplied device_map dict.
211 All-CPU and all-disk-or-CPU maps are accepted (GeneralizedComponent.__call__
212 wraps forward in Accelerate's align_module_device, so components reading raw
213 params directly still see materialized data, verified on CPU-only hardware).
214 Meta values are passed through (validated at boot against load_weights)."""
215 if isinstance(device_map, str):
216 return
217 if is_mixed_offload_gpu(device_map.values()):
218 raise ValueError(MIXED_OFFLOAD_GPU_ERROR)
221# In a mixed map, accelerate OFFLOADS the CPU/disk entries: weights live in a CPU
222# state dict or on disk, the modules hold meta placeholders, and an AlignDevicesHook
223# on the original module's forward materializes them per-call.
224# GeneralizedComponent.__call__ wraps every component call in that same hook, so this
225# is likely fine in principle — but it's only been verified on CPU-only hardware (no
226# GPU to mix in), so a map that actually puts some weights on a GPU stays rejected
227# until that's confirmed. All-CPU, all-disk, or CPU+disk maps are fine — no GPU
228# involved, so nothing to leave unverified.
229MIXED_OFFLOAD_GPU_ERROR = (
230 "device_map mixes CPU/disk offload targets with a GPU target. This is likely fine "
231 "(GeneralizedComponent.__call__ materializes offloaded params for every component "
232 "call), but has only been verified on CPU-only hardware — no GPU to mix in. Use an "
233 "all-GPU map (or n_devices) for multi-GPU, or an all-CPU/all-disk map."
234)
237def is_mixed_offload_gpu(values: Any) -> bool:
238 has_offload = has_gpu = False
239 for value in values:
240 if isinstance(value, int):
241 has_gpu = True
242 elif isinstance(value, str): 242 ↛ 239line 242 didn't jump to line 239 because the condition on line 242 was always true
243 v = value.lower()
244 if v in ("cpu", "disk"):
245 has_offload = True
246 elif v.startswith("cuda"):
247 has_gpu = True
248 return has_offload and has_gpu
251def cast_floating_params_to_dtype(model: nn.Module, dtype: torch.dtype) -> None:
252 """Cast materialized floating parameters while preserving Accelerate offload hooks.
254 Only safe on a model with no active quantizer; go through
255 ``maybe_cast_floating_params`` for anything that came out of ``from_pretrained``.
257 The one-byte-float skip below is a backstop against the worst corruption, not an
258 ownership test. Quantizer-owned scales are float32 as often as FP8: transformers'
259 finegrained-FP8 stores ``weight_scale_inv`` as float32 unless the checkpoint asks
260 for ue8m0 scales, and fbgemm-FP8 stores its scales as float32 outright. A dtype
261 cannot say who owns a tensor.
262 See: https://github.com/TransformerLensOrg/TransformerLens/issues/1743
263 """
264 from accelerate.utils import align_module_device
266 for module in model.modules():
267 with align_module_device(module):
268 for param in module.parameters(recurse=False):
269 if not param.is_floating_point() or param.dtype == dtype:
270 continue
271 if param.device.type == "meta":
272 continue
273 # Backstop, not an ownership test. One-byte floats are the case
274 # where a cast is silently unrecoverable, so they are refused even
275 # here; wider quantizer-owned scales exist and are NOT caught, which
276 # is why callers gate on the model's quantizer instead of on dtype.
277 if param.dtype.itemsize < 2:
278 continue
279 param.data = param.data.to(dtype=dtype)
282def maybe_cast_floating_params(model: nn.Module, dtype: torch.dtype) -> None:
283 """Cast floating params to dtype, skipping models with active quantization.
285 The skip is whole-model on purpose. ``from_pretrained`` has already settled the
286 load dtype by this point, and on a quantized checkpoint that is the quantizer's
287 *effective* dtype, not necessarily the requested one: quantizers may override it
288 in ``HfQuantizer.update_dtype`` (AWQ downgrades bfloat16 to float16 whenever CUDA
289 or XPU is available, whatever the placement; fbgemm-FP8 and FP-Quant force
290 bfloat16). Re-casting here would overwrite those
291 deliberate choices along with genuinely quantizer-owned storage, and dtype alone
292 cannot tell the two apart.
294 The gate releases in step with HF: a dequantized load has its
295 ``quantization_config`` deleted by ``HfQuantizer.remove_quantization_config``, so
296 ``quantization_method`` returns None and normalization resumes.
298 See: https://github.com/TransformerLensOrg/TransformerLens/issues/1713
299 See: https://github.com/TransformerLensOrg/TransformerLens/issues/1743
300 """
301 from transformer_lens.utilities.quantization import quantization_method
303 if quantization_method(getattr(model, "config", None)) is None:
304 cast_floating_params_to_dtype(model, dtype)
307def find_embedding_device(hf_model: Any) -> Optional[torch.device]:
308 """Return the device that input tokens should be placed on for a dispatched HF model.
310 When a model is loaded with ``device_map``, accelerate populates ``hf_device_map``
311 and inserts pre/post-forward hooks that route activations. Input tensors must land on
312 the device of whichever module first *consumes* them — the input embedding. Returns
313 ``None`` for single-device models (no ``hf_device_map`` set).
315 Resolves via ``hf_model.get_input_embeddings()`` rather than dict insertion order to
316 cover encoder-decoder / multimodal / audio architectures where the first entry in
317 ``hf_device_map`` is not the text-token embedding (e.g. the vision tower on LLaVA).
318 """
319 hf_device_map = getattr(hf_model, "hf_device_map", None)
320 if not hf_device_map:
321 return None
322 # Preferred: ask the model for its input embedding module and read its device.
323 get_input_embeddings = getattr(hf_model, "get_input_embeddings", None)
324 if callable(get_input_embeddings):
325 try:
326 embed_module = get_input_embeddings()
327 except (AttributeError, NotImplementedError):
328 embed_module = None
329 if embed_module is not None:
330 try:
331 param = next(embed_module.parameters())
332 return param.device
333 except StopIteration:
334 pass
335 # Fallback: first entry in hf_device_map. Less reliable but better than nothing.
336 first_device = next(iter(hf_device_map.values()))
337 if isinstance(first_device, int):
338 return torch.device("cuda", first_device)
339 return torch.device(first_device)
342def count_unique_devices(hf_model: Any) -> int:
343 """Count the number of unique devices across a dispatched HF model's ``hf_device_map``.
345 Returns 1 if the model has no ``hf_device_map`` (single-device load).
346 """
347 hf_device_map = getattr(hf_model, "hf_device_map", None)
348 if not hf_device_map:
349 return 1
350 return len(set(hf_device_map.values()))
353def find_misplaced_modules(hf_model: Any) -> list:
354 """``(module_name, mapped, actual)`` for ``hf_device_map`` entries whose loaded
355 parameters sit on a different *real* device than the map requested.
357 Accelerate places a tied parameter exactly once, so a map that splits a tie group
358 (e.g. GPT-2's ``wte``/``lm_head`` share one tensor) is silently dispatched with one
359 module's execution device pointing at weights that live elsewhere — the forward then
360 crashes deep inside a kernel. Meta parameters are skipped: they mean CPU/disk offload
361 (accelerate materializes them per-forward) or a weightless ``from_config`` load, both
362 of which are placement-consistent by construction.
363 """
364 hf_device_map = getattr(hf_model, "hf_device_map", None)
365 if not hf_device_map:
366 return []
367 misplaced = []
368 for module_name, target in hf_device_map.items():
369 try:
370 module = hf_model.get_submodule(module_name) if module_name else hf_model
371 except AttributeError:
372 continue
373 param = next(module.parameters(), None)
374 if param is None or param.device.type == "meta":
375 continue
376 expected = (
377 torch.device(f"cuda:{target}") if isinstance(target, int) else torch.device(target)
378 )
379 actual = param.device
380 same_type = actual.type == expected.type
381 same_index = expected.index is None or actual.index == expected.index
382 if not (same_type and same_index):
383 misplaced.append((module_name, str(target), str(actual)))
384 return misplaced