Coverage for transformer_lens/utilities/multi_gpu.py: 91%
145 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +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.hooked_transformer_config import (
15 HookedTransformerConfig as ConfigType,
16 )
17else:
18 ConfigType = Any
20_UNSUPPORTED_OFFLOAD_DEVICE_MAP_VALUES = {"disk"}
22AvailableDeviceMemory = list[tuple[int, int]]
23"""
24This type is passed around between different CUDA memory operations.
25The first entry of each tuple will be the device index.
26The second entry will be how much memory is currently available.
27"""
29MaxMemory = Dict[Union[str, int], Union[str, int]]
32def calculate_available_device_cuda_memory(i: int) -> int:
33 """Calculates how much memory is available at this moment for the device at the indicated index
35 Args:
36 i (int): The index we are looking at
38 Returns:
39 int: How memory is available
40 """
41 total = torch.cuda.get_device_properties(i).total_memory
42 allocated = torch.cuda.memory_allocated(i)
43 return total - allocated
46def _get_available_cuda_memory_for_device(i: int) -> int:
47 """Return a concrete byte budget accepted by Accelerate's ``max_memory``."""
48 try:
49 free_memory, _ = torch.cuda.mem_get_info(i)
50 except (AttributeError, RuntimeError):
51 return calculate_available_device_cuda_memory(i)
52 return int(free_memory)
55def determine_available_memory_for_available_devices(max_devices: int) -> AvailableDeviceMemory:
56 """Gets all available CUDA devices with their current memory calculated
58 Returns:
59 AvailableDeviceMemory: The list of all available devices with memory precalculated
60 """
61 devices = []
62 for i in range(max_devices):
63 devices.append((i, calculate_available_device_cuda_memory(i)))
65 return devices
68def sort_devices_based_on_available_memory(devices: AvailableDeviceMemory) -> AvailableDeviceMemory:
69 """Sorts all available devices with devices with the most available memory returned first
71 Args:
72 devices (AvailableDeviceMemory): All available devices with memory calculated
74 Returns:
75 AvailableDeviceMemory: The same list of passed through devices sorted with devices with most
76 available memory first
77 """
78 return sorted(devices, key=lambda x: x[1], reverse=True)
81def get_best_available_cuda_device(max_devices: Optional[int] = None) -> torch.device:
82 """Gets whichever cuda device has the most available amount of memory for use
84 Raises:
85 EnvironmentError: If there are no available devices, this will error out
87 Returns:
88 torch.device: The specific device that should be used
89 """
90 max_devices = max_devices if max_devices is not None else torch.cuda.device_count()
91 devices = determine_available_memory_for_available_devices(max_devices)
93 if len(devices) <= 0:
94 raise EnvironmentError(
95 "TransformerLens has been configured to use CUDA, but no available devices are present"
96 )
98 sorted_devices = sort_devices_based_on_available_memory(devices=devices)
100 return torch.device("cuda", sorted_devices[0][0])
103def get_best_available_device(
104 cfg: ConfigType,
105) -> torch.device:
106 """Gets the best available device to be used based on the passed in arguments
108 Args:
109 cfg: The HookedTransformerConfig object containing device configuration
111 Returns:
112 torch.device: The best available device
113 """
114 assert cfg.device is not None
115 device = torch.device(cfg.device)
117 if device.type == "cuda" and cfg.n_devices > 1: 117 ↛ 118line 117 didn't jump to line 118 because the condition on line 117 was never true
118 return get_best_available_cuda_device(cfg.n_devices)
119 else:
120 return device
123def get_device_for_block_index(
124 index: int,
125 cfg: ConfigType,
126 device: Optional[Union[torch.device, str]] = None,
127):
128 """
129 Determine the device for a given layer index based on the model configuration.
131 This function assists in distributing model layers across multiple devices. The distribution
132 is based on the configuration's number of layers (cfg.n_layers) and devices (cfg.n_devices).
135 Args:
136 index (int): Model layer index.
137 cfg: Model and device configuration.
138 device (Optional[Union[torch.device, str]], optional): Initial device used for determining the target device.
139 If not provided, the function uses the device specified in the configuration (cfg.device).
141 Returns:
142 torch.device: The device for the specified layer index.
144 Deprecated:
145 This function did not take into account a few factors for multi-GPU support. You should now
146 use get_best_available_device in order to properly run models on multiple devices.
147 This will be removed in 3.0
148 """
149 assert cfg.device is not None
150 if device is None:
151 device = cfg.device
152 device = torch.device(device)
153 if device.type == "cpu":
154 return device
155 # Multiplying first guarantees the result is in [0, n_devices - 1] and avoids
156 # the divide-by-zero when n_layers < n_devices. The naive form
157 # `index // (n_layers // n_devices)` floors the divisor and overshoots when
158 # n_layers is not a multiple of n_devices (e.g. 62 layers / 8 devices → 8).
159 device_index = (device.index or 0) + (index * cfg.n_devices) // cfg.n_layers
160 return torch.device(device.type, device_index)
163def resolve_device_map(
164 n_devices: Optional[int],
165 device_map: Optional[Union[str, Dict[str, Union[str, int]]]],
166 device: Optional[Union[str, torch.device]],
167 max_memory: Optional[MaxMemory] = None,
168) -> Tuple[Optional[Union[str, Dict[str, Union[str, int]]]], Optional[MaxMemory]]:
169 """Resolve ``n_devices`` / ``device_map`` / ``device`` into HF ``from_pretrained`` kwargs.
171 Returns ``(device_map, max_memory)`` tuple ready to pass into ``model_kwargs``.
173 Semantics:
174 - Explicit ``device_map`` wins and is passed through unchanged (user-provided
175 ``max_memory`` is passed through too). CPU targets are supported; disk / meta
176 offload targets are still rejected because Bridge component wrappers can bypass
177 Accelerate's offload hooks during forward passes.
178 - ``n_devices=None`` or ``1``: returns ``(None, None)`` — single-device path.
179 - ``n_devices > 1``: returns ``("balanced", {0: bytes, ..., n-1: bytes})``.
180 ``"balanced"`` is accelerate's string directive for balanced layer dispatch;
181 concrete byte budgets cap visibility to exactly ``n_devices`` GPUs.
182 """
183 if device_map is not None and device is not None:
184 raise ValueError("device and device_map are mutually exclusive — pass one.")
185 if device_map is not None:
186 _validate_device_map_values(device_map)
187 return device_map, max_memory
188 if n_devices is None or n_devices <= 1:
189 return None, max_memory
190 if not torch.cuda.is_available():
191 raise ValueError(f"n_devices={n_devices} requires CUDA, which is not available.")
192 if torch.cuda.device_count() < n_devices: 192 ↛ 193line 192 didn't jump to line 193 because the condition on line 192 was never true
193 raise ValueError(
194 f"n_devices={n_devices} but only {torch.cuda.device_count()} CUDA devices present."
195 )
196 resolved_max_memory: MaxMemory = (
197 dict(max_memory)
198 if max_memory is not None
199 else {i: _get_available_cuda_memory_for_device(i) for i in range(n_devices)}
200 )
201 return "balanced", resolved_max_memory
204def _validate_device_map_values(
205 device_map: Union[str, Dict[str, Union[str, int]]],
206) -> None:
207 """Reject explicit disk values and mixed CPU+GPU targets in a user-supplied
208 device_map dict. Meta values are passed through (validated at boot against
209 load_weights)."""
210 if isinstance(device_map, str):
211 return
212 for key, value in device_map.items():
213 normalized = str(value).lower() if isinstance(value, str) else None
214 if normalized in _UNSUPPORTED_OFFLOAD_DEVICE_MAP_VALUES:
215 raise ValueError(
216 f"device_map[{key!r}]={value!r} is not supported yet. TransformerBridge "
217 "currently supports CPU device_map targets, but disk / meta offload can "
218 "bypass Accelerate hooks inside wrapped Bridge components."
219 )
220 if is_mixed_cpu_gpu(device_map.values()):
221 raise ValueError(MIXED_CPU_GPU_ERROR)
224# In a mixed map, accelerate OFFLOADS the CPU entries: weights live in a CPU state
225# dict, the modules hold meta placeholders, and an AlignDevicesHook on the original
226# module's forward materializes them per-call. Bridge components that compute from raw
227# parameters (e.g. NormalizationBridge reads self.weight to expose hook_normalized)
228# never trigger that hook, so the forward hits meta tensors. All-CPU maps are fine —
229# no offload, real parameters.
230MIXED_CPU_GPU_ERROR = (
231 "device_map mixes CPU and GPU targets, which accelerate implements as CPU offload "
232 "(meta placeholders materialized by forward hooks that Bridge components bypass). "
233 "Use an all-GPU map (or n_devices) for multi-GPU, or an all-CPU map."
234)
237def is_mixed_cpu_gpu(values: Any) -> bool:
238 has_cpu = 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 == "cpu":
245 has_cpu = True
246 elif v.startswith("cuda"):
247 has_gpu = True
248 return has_cpu 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 Skips one-byte floats (FP8 dtypes like float8_e8m0fnu) which are quantizer-owned
255 scale parameters — casting them corrupts the quantization format.
256 """
257 from accelerate.utils import align_module_device
259 for module in model.modules():
260 with align_module_device(module):
261 for param in module.parameters(recurse=False):
262 if not param.is_floating_point() or param.dtype == dtype:
263 continue
264 if param.device.type == "meta":
265 continue
266 # Skip one-byte floats (FP8 scale tensors): they are quantizer-owned
267 # and casting them breaks the weight/scale pair relationship.
268 if param.dtype.itemsize < 2:
269 continue
270 param.data = param.data.to(dtype=dtype)
273def maybe_cast_floating_params(model: nn.Module, dtype: torch.dtype) -> None:
274 """Cast floating params to dtype, skipping models with active quantization.
276 When a model has an active quantization_config, the quantizer owns specific
277 dtypes (e.g., FP8 scales) that must not be overwritten. This helper wraps
278 the cast with that check.
280 See: https://github.com/TransformerLensOrg/TransformerLens/issues/1713
281 """
282 from transformer_lens.utilities.quantization import quantization_method
284 if quantization_method(getattr(model, "config", None)) is None:
285 cast_floating_params_to_dtype(model, dtype)
288def find_embedding_device(hf_model: Any) -> Optional[torch.device]:
289 """Return the device that input tokens should be placed on for a dispatched HF model.
291 When a model is loaded with ``device_map``, accelerate populates ``hf_device_map``
292 and inserts pre/post-forward hooks that route activations. Input tensors must land on
293 the device of whichever module first *consumes* them — the input embedding. Returns
294 ``None`` for single-device models (no ``hf_device_map`` set).
296 Resolves via ``hf_model.get_input_embeddings()`` rather than dict insertion order to
297 cover encoder-decoder / multimodal / audio architectures where the first entry in
298 ``hf_device_map`` is not the text-token embedding (e.g. the vision tower on LLaVA).
299 """
300 hf_device_map = getattr(hf_model, "hf_device_map", None)
301 if not hf_device_map:
302 return None
303 # Preferred: ask the model for its input embedding module and read its device.
304 get_input_embeddings = getattr(hf_model, "get_input_embeddings", None)
305 if callable(get_input_embeddings):
306 try:
307 embed_module = get_input_embeddings()
308 except (AttributeError, NotImplementedError):
309 embed_module = None
310 if embed_module is not None:
311 try:
312 param = next(embed_module.parameters())
313 return param.device
314 except StopIteration:
315 pass
316 # Fallback: first entry in hf_device_map. Less reliable but better than nothing.
317 first_device = next(iter(hf_device_map.values()))
318 if isinstance(first_device, int):
319 return torch.device("cuda", first_device)
320 return torch.device(first_device)
323def count_unique_devices(hf_model: Any) -> int:
324 """Count the number of unique devices across a dispatched HF model's ``hf_device_map``.
326 Returns 1 if the model has no ``hf_device_map`` (single-device load).
327 """
328 hf_device_map = getattr(hf_model, "hf_device_map", None)
329 if not hf_device_map:
330 return 1
331 return len(set(hf_device_map.values()))
334def find_misplaced_modules(hf_model: Any) -> list:
335 """``(module_name, mapped, actual)`` for ``hf_device_map`` entries whose loaded
336 parameters sit on a different *real* device than the map requested.
338 Accelerate places a tied parameter exactly once, so a map that splits a tie group
339 (e.g. GPT-2's ``wte``/``lm_head`` share one tensor) is silently dispatched with one
340 module's execution device pointing at weights that live elsewhere — the forward then
341 crashes deep inside a kernel. Meta parameters are skipped: they mean CPU/disk offload
342 (accelerate materializes them per-forward) or a weightless ``from_config`` load, both
343 of which are placement-consistent by construction.
344 """
345 hf_device_map = getattr(hf_model, "hf_device_map", None)
346 if not hf_device_map:
347 return []
348 misplaced = []
349 for module_name, target in hf_device_map.items():
350 try:
351 module = hf_model.get_submodule(module_name) if module_name else hf_model
352 except AttributeError:
353 continue
354 param = next(module.parameters(), None)
355 if param is None or param.device.type == "meta":
356 continue
357 expected = (
358 torch.device(f"cuda:{target}") if isinstance(target, int) else torch.device(target)
359 )
360 actual = param.device
361 same_type = actual.type == expected.type
362 same_index = expected.index is None or actual.index == expected.index
363 if not (same_type and same_index):
364 misplaced.append((module_name, str(target), str(actual)))
365 return misplaced