base.py 70.3 KB
Newer Older
1
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
4
5
#
# See LICENSE for license information.

"""Base modules and utilities for TransformerEngine PyTorch API"""
6
import io
7
import math
8
9
10
import os
import pickle
import warnings
11
from enum import Enum
12
from abc import ABC, abstractmethod
13
from typing import Any, Dict, Generator, List, Optional, Tuple, Union
14
from contextlib import contextmanager
15
from types import MethodType
16
17
18

import torch
import torch.nn.functional as F
19
from torch.distributed.tensor import DTensor
20

21
import transformer_engine_torch as tex
22

23
from ._common import _ParameterInitMeta, noop_cat
24
from ..quantization import (
25
26
    MXFP8BlockScalingRecipeState,
    DelayedScalingRecipeState,
27
    Float8CurrentScalingRecipeState,
28
    Float8BlockScalingRecipeState,
29
    NVFP4BlockScalingRecipeState,
30
    FP8GlobalStateManager,
31
    RecipeState,
32
33
34
35
36
)
from ..distributed import (
    gather_along_first_dim,
    is_fp8_activation_recompute_enabled,
    in_fp8_activation_recompute_phase,
37
    _fsdp_gather_tensors,
38
39
)
from ..constants import dist_group_type
40
from ..cpp_extensions.gemm import _NUM_MAX_UB_STREAMS
41
from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer
42
43
from ..tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer
from ..tensor.mxfp8_tensor import MXFP8Quantizer
44
from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer
45
46
from ..tensor.storage.float8_tensor_storage import Float8TensorStorage
from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage
47
from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage
48
49
50
51
from ..utils import (
    is_non_tn_fp8_gemm_supported,
    torch_get_autocast_gpu_dtype,
    get_nvtx_range_context,
52
53
    nvtx_range_push,
    nvtx_range_pop,
54
)
55
from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage
56
from ...common.recipe import DelayedScaling, Recipe
57
58
from ...debug.pytorch.debug_state import TEDebugState
from ...debug.pytorch.debug_quantization import DebugQuantizer, DebugQuantizedTensor
59
from ...debug.pytorch.utils import next_iter_when_debug_should_be_run, any_feature_enabled
60

61
__all__ = ["initialize_ub", "destroy_ub", "UserBufferQuantizationMode"]
62

63
64
65
_2X_ACC_FPROP = False
_2X_ACC_DGRAD = True
_2X_ACC_WGRAD = True
66
_dummy_wgrads = {}
67
_ub_communicators = None
68
_MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY = None, None
69
layers_atomic_ring_exchange = []
70
71


72
73
74
75
76
77
78
79
80
class UserBufferQuantizationMode(Enum):
    """
    UserBufferQuantizationMode is an enum that represents the quantization mode of the UserBuffer.
    """

    NONE = "none"
    FP8 = "fp8"


81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def get_dummy_wgrad(shape: list, dtype: torch.dtype, zero=False) -> torch.Tensor:
    """Returns a dummy tensor of given shape."""
    assert len(shape) == 2
    global _dummy_wgrads
    if (shape[0], shape[1], dtype) not in _dummy_wgrads:
        _dummy_wgrads[(shape[0], shape[1], dtype)] = torch.empty(
            shape,
            dtype=dtype,
            device="cuda",
            requires_grad=False,
        )
    if zero:
        _dummy_wgrads[(shape[0], shape[1], dtype)].fill_(0)
    return _dummy_wgrads[(shape[0], shape[1], dtype)].detach()


97
98
def initialize_ub(
    shape: list,
99
    tp_size: int,
100
    use_fp8: bool = False,
101
    quantization_modes: List[UserBufferQuantizationMode] = None,
102
    dtype: torch.dtype = torch.bfloat16,
103
    ub_cfgs: Optional[Union[dict, List[dict]]] = None,
104
    bootstrap_backend: Union[str, torch.distributed.Backend] = None,
105
) -> None:
106
107
    r"""
    Initialize the Userbuffers communicator for overlapping tensor-parallel communications with
Paweł Gadziński's avatar
Paweł Gadziński committed
108
    GEMM compute in ``te.Linear``, ``te.LayerNormLinear`` and ``te.LayerNormMLP`` modules.
109
110
111
112
113

    Parameters
    ----------
    shape : list
            shape of the communication buffer, typically set to be the same as the global shape of
Paweł Gadziński's avatar
Paweł Gadziński committed
114
115
            the input tensor to a ``te.TransformerLayer`` forward pass, with the sequence and batch
            dimensions collapsed together -- i.e.: ``(sequence_length * batch_size, hidden_size)``
116
117
118
    tp_size : int
              number of GPUs in the tensor-parallel process group
    use_fp8 : bool = False
119
              allocate the communication buffer for FP8 GEMM inputs/outputs.
Paweł Gadziński's avatar
Paweł Gadziński committed
120
              DEPRECATED: Please use ``quantization_modes`` instead.
121
122
    quantization_modes : List[UserBufferQuantizationMode] = None
              if a list of UserBufferQuantizationMode is provided, a UB communicator is created for each quantization setting in the list.
Paweł Gadziński's avatar
Paweł Gadziński committed
123
              falls back to the legacy ``use_fp8`` parameter if ``None`` is provided.
124
    dtype : torch.dtype = torch.bfloat16
Paweł Gadziński's avatar
Paweł Gadziński committed
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
            non-FP8 data type of the communication buffer when ``use_fp8 = False``
    ub_cfgs : dict = None
             Configuration dictionary with the structure::

                 {
                    <gemm_name> : {
                        "method": <"ring_exchange" or "pipeline">,
                        "is_reduce_scatter": bool,
                        "num_sm": int,
                        "cga_size": int,
                        "set_sm_margin": bool,
                        "num_splits": int,
                        "aggregate": bool,
                        "atomic_gemm": bool,
                        "use_ce": bool,
                        "fp8_buf": bool,
                    }
                 }

             for ``te.TransformerLayer`` GEMM layers in ``["qkv_fprop", "qkv_dgrad", "qkv_wgrad",
145
             "proj_fprop", "proj_dgrad", "proj_wgrad", "fc1_fprop", "fc1_dgrad", "fc2_dgrad",
Paweł Gadziński's avatar
Paweł Gadziński committed
146
147
             "fc2_fprop", "fc2_wgrad"]``.
             a list may be provided to specify different overlap configurations for different the quantization settings in ``quantization_modes``
148
    bootstrap_backend : str = None
Paweł Gadziński's avatar
Paweł Gadziński committed
149
                        ``torch.distributed`` communication backend for the all-gather, broadcast and
150
151
152
153
                        barrier collectives during Userbuffers initialization. Not all backends are
                        valid for every cluster configuration and distributed launch method even if
                        they are available in PyTorch. When left unset, the initialization prefers
                        to use the MPI backend, falling back first on Gloo and then NCCL if MPI is
Paweł Gadziński's avatar
Paweł Gadziński committed
154
                        not available. Setting ``NVTE_UB_WITH_MPI=1`` when building TE overrides this
155
                        option and always initializes Userbuffers with direct MPI calls in C++,
Paweł Gadziński's avatar
Paweł Gadziński committed
156
                        which also requires ``MPI_HOME=/path/to/mpi/root`` to be set at compile time.
157
    """
158
    if not tex.device_supports_multicast():
159
        assert bool(int(os.getenv("UB_SKIPMC", "0"))), (
160
161
162
163
            "CUDA device, driver and/or toolkit version does not support comm+GEMM overlap with "
            + "CUDA Multicast. Launch app with UB_SKIPMC=1 to try CUDA IPC instead."
        )

164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
    if not quantization_modes:
        warnings.warn(
            "Initializing Userbuffers with use_fp8 is deprecated. Please use quantization_modes"
            " instead.",
            DeprecationWarning,
        )
        quantization_modes = [
            UserBufferQuantizationMode.FP8 if use_fp8 else UserBufferQuantizationMode.NONE
        ]
    else:
        assert isinstance(quantization_modes, list), "quantization_modes must be a list"
        assert all(
            isinstance(mode, UserBufferQuantizationMode) for mode in quantization_modes
        ), "quantization_modes must be a list of UserBufferQuantizationMode"

    if isinstance(ub_cfgs, dict) or ub_cfgs is None:
        ub_cfgs = [ub_cfgs] * len(quantization_modes)
    else:
        assert len(ub_cfgs) == len(
            quantization_modes
        ), "Number of ub_cfgs settings must match number of quantization configurations"

186
187
188
    global _ub_communicators
    assert _ub_communicators is None, "UB communicators are already initialized."
    _ub_communicators = {}
189
190

    if tex.ubuf_built_with_mpi():
191
192
        # We're bootstrapping with direct calls to MPI in Userbuffers code so we need to force
        # an MPI_Init() here by creating a new MPI process group...
193
        assert torch.distributed.is_mpi_available()
194
195
        _ = torch.distributed.new_group(backend="mpi")
        helper = tex.CommOverlapHelper()
196
    else:
197
198
        # Bootstrapping with torch.distributed API, so check backend and construct
        # intra/inter-node process groups...
199
200
201
202
203
        assert (
            torch.distributed.is_initialized()
        ), "torch.distributed must be initialized before Userbuffers"
        if bootstrap_backend is None:
            bootstrap_backend = "nccl"
204
            if torch.distributed.is_mpi_available():
205
                bootstrap_backend = "mpi"
206
207
            elif torch.distributed.is_gloo_available():
                bootstrap_backend = "gloo"
208
        else:
209
210
211
212
213
214
215
216
217
            assert bootstrap_backend in [
                "gloo",
                "mpi",
                "nccl",
            ], "Invalid torch.distributed backend for bootstrapping Userbuffers!"
            assert torch.distributed.is_backend_available(bootstrap_backend), (
                f"PyTorch must be compiled with '{bootstrap_backend}' support in order to "
                f"bootstrap Userbuffers with '{bootstrap_backend}' collectives."
            )
218
219
220
221
222

        world_group = torch.distributed.new_group(backend=bootstrap_backend)
        world_rank = torch.distributed.get_rank(world_group)
        world_size = torch.distributed.get_world_size(world_group)

223
224
        num_domains = world_size // tp_size
        mydomain_idx = world_rank // tp_size
225
        if num_domains > 1:
226
227
228
229
            ranks_per_domain_list = [
                [i * tp_size + t for t in range(tp_size)] for i in range(num_domains)
            ]
            tp_domain_group, _ = torch.distributed.new_subgroups_by_enumeration(
230
231
                ranks_per_domain_list, backend=bootstrap_backend
            )
232
233
            local_rank = torch.distributed.get_rank(tp_domain_group)
            tp_domain_ranks = torch.distributed.get_process_group_ranks(tp_domain_group)
234

235
            helper = tex.CommOverlapHelper(world_group, tp_domain_group)
236
        else:
237
238
            # TP model on single NVLink domain, no replication, no data-parallelism
            mydomain_idx = 0
239
            local_rank = world_rank
240
            tp_domain_ranks = list(range(world_size))
241
242

            helper = tex.CommOverlapHelper(world_group)
243

244
        if world_rank == 0:
245
            print(f"!!! [UB] Number of TP domains: {num_domains}\n", end="", flush=True)
246
247
        if local_rank == 0:
            print(
248
                f"!!! [UB] Global ranks on TP domain {mydomain_idx}: {tp_domain_ranks}\n",
249
250
251
252
                end="",
                flush=True,
            )

253
    # Default buffer precision: AllGather buffers use fp8 when using fp8 recipe
254
    layers_all_gather_overlap = [
255
256
257
        "qkv_fprop",
        "qkv_dgrad",
        "proj_dgrad",
258
        "proj_wgrad",
259
260
261
        "fc1_fprop",
        "fc1_dgrad",
        "fc2_dgrad",
262
        "fc2_wgrad",
263
    ]
264
    layers_reduce_scatter_overlap = ["proj_fprop", "fc2_fprop", "qkv_wgrad", "fc1_wgrad"]
Jaemin Choi's avatar
Jaemin Choi committed
265
    dgrad_reduce_scatter_overlap = ["qkv_dgrad", "fc1_dgrad"]
266
267
    # Default overlap methods for layers
    methods = {
268
269
270
271
272
273
        "ring_exchange": [
            "qkv_fprop",
            "fc1_fprop",
            "proj_dgrad",
            "fc2_dgrad",
        ],
274
275
        "pipeline": ["proj_fprop", "fc2_fprop"],
        "bulk": ["qkv_dgrad", "qkv_wgrad", "fc1_dgrad", "fc1_wgrad"],
276
        "external": ["proj_wgrad", "fc2_wgrad"],
277
278
    }

279
    # AG-RS overlap pairs of layers forming a tensor-parallel block
280
281
    ag_rs_pairs = {"qkv_fprop": "proj_fprop", "fc1_fprop": "fc2_fprop"}
    rs_ag_pairs = {v: k for k, v in ag_rs_pairs.items()}
282
    external_gemm_to_overlap = {"proj_wgrad": "proj_dgrad", "fc2_wgrad": "fc2_dgrad"}
283
284
285
    global layers_atomic_ring_exchange
    layers_atomic_ring_exchange = []

286
287
288
289
290
291
    def get_method(name):
        for method, names in methods.items():
            if name in names:
                return method
        raise KeyError(f"Given layer name {name} does not exist.")

292
    def get_default_config(name):
293
        global _MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY
294
295
        method = get_method(name)
        is_reduce_scatter = name in layers_reduce_scatter_overlap
296
297
        if _MIN_STREAM_PRIORITY is None or _MAX_STREAM_PRIORITY is None:
            _MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY = tex.get_stream_priority_range()
298
299
300
301
302
        default_cfg = {
            "method": method,
            "is_reduce_scatter": is_reduce_scatter,
            "num_sm": 1 if method == "ring_exchange" else 16,
            "cga_size": 1 if method == "ring_exchange" else 2,
303
304
            "set_sm_margin": not method == "ring_exchange",
            "num_splits": tp_size if method == "ring_exchange" else 4,
305
306
307
308
            "aggregate": False,
            "atomic_gemm": False,
            "use_ce": True,
            "fp8_buf": name in layers_all_gather_overlap,
309
310
311
            "comm_priority": _MAX_STREAM_PRIORITY,
            "gemm_priority": _MIN_STREAM_PRIORITY,
            "pipeline_rs_overlap_first_gemm": False,
312
313
314
        }
        return default_cfg

315
316
    def add_ub(
        name: str,
317
        quantization_mode: UserBufferQuantizationMode,
318
        method: str,
319
        is_reduce_scatter: bool,
320
321
        num_sm: int = 16,
        cga_size: int = 2,
322
        set_sm_margin: bool = False,
323
        num_splits: int = 0,
324
325
        aggregate: bool = False,
        atomic_gemm: bool = False,
326
        use_ce: bool = True,
327
        fp8_buf: bool = False,
328
329
330
        comm_priority: int = 0,
        gemm_priority: int = 0,
        pipeline_rs_overlap_first_gemm: bool = False,
331
    ) -> None:
332
333
334
335
        if atomic_gemm:
            warnings.warn(
                "Atomic GEMM uses a beta API from cublas and is not tested for all use cases."
            )
336
337
338
            assert (
                quantization_mode == UserBufferQuantizationMode.FP8
            ), "Atomic GEMM overlap supported only for FP8 GEMM."
339
            if method in ("bulk", "external"):
340
                warnings.warn(
341
                    f"At {name}, atoimic GEMM not is supported for a bulk overlap."
342
343
344
                    "Defaulting to `atomic_gemm=False`."
                )
                atomic_gemm = 0
345
        if not is_reduce_scatter and method == "pipeline":
346
            raise ValueError(
347
                f"At {name}, `pipeline` overlap method is not supported for AllGather."
348
            )
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
        # Check if both AG and RS overlaps use `atomic GEMM`` + `p2p ring-exchange`.
        # Using atomic GEMM + p2p ring-exchange in only one of the pair breaks functionality.
        global layers_atomic_ring_exchange
        if atomic_gemm and method == "ring_exchange" and name in ag_rs_pairs:
            layers_atomic_ring_exchange += [name, ag_rs_pairs[name]]
        if name in rs_ag_pairs:
            assert_message = (
                f"At {name}, atomic AG-GEMM overlap with `ring_exchange` shuffles GEMM chunk "
                "outputs, and  RS-GEMM overlap un-suffle them. When one of the GEMM-AG and "
                "GEMM-RS overlaps forming a TP block (e.g., qkv_fprop and proj_fprop) uses "
                "`atomic gemm` and `ring_exhcnage`, its pair must use the same overlap config "
                "for functionality."
            )
            if name in layers_atomic_ring_exchange:
                assert atomic_gemm and method == "ring_exchange", assert_message
            else:
                if atomic_gemm and method == "ring_exchange":
                    assert rs_ag_pairs[name] in layers_atomic_ring_exchange, assert_message

368
369
370
371
372
373
374
375
376
377
        if name in external_gemm_to_overlap:
            assert method == "external", (
                f"At {name}, `external` overlap method is specified, but the selected method is"
                f" {method}"
            )
            assert external_gemm_to_overlap[name] in methods["ring_exchange"], (
                f"At {name}, `external` overlap method is specified, but the external gemm"
                f" {external_gemm_to_overlap[name]} is not using `ring_exchange` overlap method"
            )

378
379
380
381
382
        buffer_dtype = (
            torch.uint8
            if (quantization_mode == UserBufferQuantizationMode.FP8 and fp8_buf)
            else dtype
        )
383
        if method == "ring_exchange":
384
385
386
387
            ub_obj = tex.CommOverlapP2P(
                shape,  # Communication buffer shape
                buffer_dtype,  # Communication buffer data type
                helper,  # Helper for torch.distributed callbacks during bootstrapping
388
                tp_size,  # Tensor-parallel group size (may be different than local_size)
389
390
391
392
393
394
395
396
                tex.CommOverlapType.RS if is_reduce_scatter else tex.CommOverlapType.AG,
                num_max_streams=_NUM_MAX_UB_STREAMS,
                comm_cga_size=cga_size,
                num_comm_sm=num_sm,
                set_sm_margin=set_sm_margin,
                atomic_gemm=atomic_gemm,
                use_ce=use_ce,
                aggregate=aggregate,
397
398
                gemm_priority=gemm_priority,
                comm_priority=comm_priority,
399
            )
400
        else:
401
402
403
404
            ub_obj = tex.CommOverlap(
                shape,  # Communication buffer shape
                buffer_dtype,  # Communication buffer data type
                helper,  # Helper for torch.distributed callbacks during bootstrapping
405
                tp_size,  # Tensor-parallel group size (may be different than local_size)
406
407
408
409
410
411
                num_splits=num_splits,
                num_max_streams=_NUM_MAX_UB_STREAMS,
                comm_cga_size=cga_size,
                num_comm_sm=num_sm,
                set_sm_margin=set_sm_margin,
                atomic_gemm=atomic_gemm,
412
413
414
                gemm_priority=gemm_priority,
                comm_priority=comm_priority,
                rs_overlap_first_gemm=pipeline_rs_overlap_first_gemm,
415
            )
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
        _ub_communicators[(name, quantization_mode)] = ub_obj

    for quantization_mode, user_ub_cfg in zip(quantization_modes, ub_cfgs):
        if user_ub_cfg is not None:
            for name in dgrad_reduce_scatter_overlap:
                if (
                    name in user_ub_cfg
                    and "method" in user_ub_cfg[name]
                    and user_ub_cfg[name]["method"] != "bulk"
                ):
                    wgrad_name = name.replace("dgrad", "wgrad")
                    assert wgrad_name not in user_ub_cfg
                    layers_reduce_scatter_overlap.remove(wgrad_name)
                    layers_all_gather_overlap.remove(name)
                    layers_reduce_scatter_overlap.append(name)
                    methods["bulk"].remove(name)
                    new_method = user_ub_cfg[name]["method"]
                    methods[new_method].append(name)

        for name in (
            methods["ring_exchange"] + methods["pipeline"] + methods["bulk"] + methods["external"]
        ):
            ub_cfg = get_default_config(name)
            if user_ub_cfg is not None and name in user_ub_cfg:
                fp8_buf = (name in layers_all_gather_overlap) or (
                    user_ub_cfg[name].get("fp8_buf", False) and name in methods["pipeline"]
                )
443
                ub_cfg.update(user_ub_cfg[name])
444
445
                ub_cfg["fp8_buf"] = fp8_buf
            add_ub(name, quantization_mode, **ub_cfg)
446
447


448
def get_ub(name: str, use_fp8: bool):
449
    """Get userbuffer communicator corresponding to give key."""
450
451
452
453
    # For now use `use_fp8` boolean input as it matches the current design in the modules
    # So favour simplicity until the correct design becomes clear.
    # This is mainly an internal API so we don't need to worry about future changes
    key = (name, UserBufferQuantizationMode.FP8 if use_fp8 else UserBufferQuantizationMode.NONE)
454
    assert _ub_communicators is not None, "UB manager is not initialized."
455
456
    assert key in _ub_communicators, f"UB for {name} with use_fp8={use_fp8} is not registered."
    return _ub_communicators[key]
457

458

459
460
461
462
463
464
465
def destroy_ub():
    """Destroy all allocated userbuffer communicators."""
    global _ub_communicators
    _ub_communicators = None
    global layers_atomic_ring_exchange
    layers_atomic_ring_exchange = []

466

467
468
469
470
471
def fill_userbuffers_buffer_for_all_gather(
    comm,
    local_tensor: torch.Tensor,
    quantizer: Optional[Quantizer],
    process_group,
472
) -> tuple[torch.Tensor | QuantizedTensorStorage, torch.Tensor | QuantizedTensorStorage]:
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
    """Fill local shard of Userbuffers buffer with data for all-gather

    Returns the full tensor and the local shard, both using the
    Userbuffers buffer as their underlying data. These tensors should
    be used carefully (e.g. only immediately before and after a
    Userbuffers operation) since the underlying data may be
    overwritten by other Userbuffers operations.

    May perform blocking communication if needed for the gathered
    tensor's metadata, e.g. scaling factors.

    """

    # Tensor dimensions
    local_shape = local_tensor.size()
    if not local_shape:
        raise ValueError(f"Invalid local tensor (shape={tuple(local_shape)})")
    process_group_size = torch.distributed.get_world_size(process_group)
    global_shape = list(local_shape)
    global_shape[0] *= process_group_size

    # Unquantized data
    if quantizer is None:
496
        if isinstance(local_tensor, QuantizedTensorStorage):
497
498
499
500
501
502
503
504
505
506
507
508
            local_tensor = local_tensor.dequantize()
        if comm.is_fp8_ubuf():
            raise RuntimeError(
                "Attempting to all-gather unquantized tensor, "
                "but Userbuffers is initialized with FP8 buffers"
            )
        comm.copy_into_buffer(local_tensor, local_chunk=True)
        global_tensor = comm.get_buffer(shape=global_shape)
        return global_tensor, local_tensor

    # FP8 data
    if isinstance(quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer)):
509
510
        if not isinstance(local_tensor, Float8TensorStorage):
            if isinstance(local_tensor, QuantizedTensorStorage):
511
512
513
514
515
516
517
518
519
520
                local_tensor.dequantize()
            quantizer.set_usage(rowwise=True, columnwise=False)
            local_tensor = quantizer(local_tensor)
        if not comm.is_fp8_ubuf():
            raise RuntimeError(
                "Attempting to all-gather FP8 tensor, "
                "but Userbuffers is not initialized with FP8 buffers"
            )
        comm.copy_into_buffer(local_tensor._data, local_chunk=True)
        global_tensor_data = comm.get_buffer(shape=global_shape)
521
        global_tensor = Float8TensorStorage(
522
523
524
525
526
527
528
529
530
531
532
            data=global_tensor_data,
            fp8_scale_inv=local_tensor._scale_inv,
            fp8_dtype=local_tensor._fp8_dtype,
            quantizer=quantizer,
        )
        return global_tensor, local_tensor

    # MXFP8 data
    if isinstance(quantizer, MXFP8Quantizer):

        # Cast to MXFP8 if needed
533
534
        if not isinstance(local_tensor, MXFP8TensorStorage):
            if isinstance(local_tensor, QuantizedTensorStorage):
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
                local_tensor.dequantize()
            local_tensor = quantizer(local_tensor)
        if not comm.is_fp8_ubuf():
            raise RuntimeError(
                "Attempting to all-gather MXFP8 tensor, "
                "but Userbuffers is not initialized with FP8 buffers"
            )

        # Check which MXFP8 buffer to communicate
        if quantizer.rowwise_usage == quantizer.columnwise_usage:
            raise ValueError(
                "Userbuffers can only communicate one MXFP8 buffer at a time, "
                f"but quantizer has rowwise_usage={quantizer.rowwise_usage}, "
                f"columnwise_usage={quantizer.columnwise_usage}"
            )
        with_rowwise_data = quantizer.rowwise_usage

        # Copy MXFP8 data to local chunk of Userbuffers buffer
        local_data = (
            local_tensor._rowwise_data if with_rowwise_data else local_tensor._columnwise_data
        )
        comm.copy_into_buffer(local_data, local_chunk=True)

        # Gather scaling-inverses
        if math.prod(local_shape[:-1]) % 128 != 0:
            raise ValueError(
                "Userbuffers requires MXFP8 tensor dims that are divisible by 128, "
                f"but got MXFP8 tensor with shape={tuple(local_shape)}"
            )
564
565
        if local_tensor._with_gemm_swizzled_scales:
            raise ValueError("Userbuffers assumes MXFP8 tensors have unswizzled scales")
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
        local_scale_inv = (
            local_tensor._rowwise_scale_inv
            if with_rowwise_data
            else local_tensor._columnwise_scale_inv
        )
        local_scale_inv_size = list(local_scale_inv.size())
        global_scale_inv = torch.empty(
            [process_group_size * local_scale_inv_size[0]] + local_scale_inv_size[1:],
            dtype=local_scale_inv.dtype,
            device=local_scale_inv.device,
        )
        torch.distributed.all_gather_into_tensor(
            global_scale_inv,
            local_scale_inv,
            group=process_group,
        )

        # Construct MXFP8 tensor with Userbuffers buffer
        rowwise_data, rowwise_scale_inv = None, None
        columnwise_data, columnwise_scale_inv = None, None
        global_data = comm.get_buffer(shape=global_shape)
        if with_rowwise_data:
            rowwise_data, rowwise_scale_inv = global_data, global_scale_inv
        else:
            columnwise_data, columnwise_scale_inv = global_data, global_scale_inv
591
        global_tensor = MXFP8TensorStorage(
592
593
594
595
596
597
            rowwise_data=rowwise_data,
            rowwise_scale_inv=rowwise_scale_inv,
            columnwise_data=columnwise_data,
            columnwise_scale_inv=columnwise_scale_inv,
            fp8_dtype=local_tensor._fp8_dtype,
            quantizer=quantizer,
598
            with_gemm_swizzled_scales=False,
599
600
601
602
603
604
605
        )
        return global_tensor, local_tensor

    # Unsupported data format
    raise ValueError(f"Unsupported quantizer for Userbuffers ({quantizer})")


606
607
608
class TransformerEngineBaseModule(torch.nn.Module, ABC):
    """Base TE module."""

609
    def __init__(self, name: Optional[str] = None) -> None:
610
611
        super().__init__()
        assert torch.cuda.is_available(), "TransformerEngine needs CUDA."
612
        self.name = name
613
        self.next_iter_when_debug_should_be_run = 0
614
615
616
617
        self.fp8_initialized = False
        self.fp8 = False
        self.fp8_calibration = False
        self.fp8_meta = {}
618
        self.fp8_meta["fp8_checkpoint"] = False
619
620
        self.fp8_meta["fp8_group"] = None
        self.fp8_meta_tensors_initialized = False
621
        self.quantizers = {"scaling_fwd": {}, "scaling_bwd": {}}
622
623
624
        self.tp_group = None
        self.tp_size = 1
        self.sequence_parallel = False
625
626
        self.param_init_meta = {}
        self.primary_weights_in_fp8 = FP8GlobalStateManager.with_fp8_parameters()
627
        self.preserve_high_precision_init_val = FP8GlobalStateManager.with_high_precision_init_val()
628
629
        self.fsdp_wrapped = False
        self.fsdp_group = None
630
        self._fp8_workspaces: Dict[str, QuantizedTensor] = {}
631
        self.activation_dtype: Optional[torch.dtype] = None
632
        self.wgrad_accumulation_and_reduce_hooks = []
633
        self.wgrad_store = None
634

635
636
        if not TEDebugState.debug_enabled:
            TEDebugState.initialize()
637
        self._validate_name()
638

639
640
641
642
643
644
    def fast_setattr(self, name: str, value: Any) -> None:
        """
        Fast version of the Module's set attribute function.
        Should be used for regular attributes, but not properties nor parameters/buffers.
        """
        self.__dict__[name] = value
645

646
647
648
649
650
651
652
    def module_setattr(self, name: str, value: Any) -> None:
        """
        Regular version of the Module's set attribute function.
        Should be used only when the fast version cannot be used - for the properties,
        parameters and buffers.
        """
        super().__setattr__(name, value)
653

654
    def adjust_amax_history_length(self, length: int, fwd: Optional[bool] = None) -> None:
655
656
657
658
        """
        Delayed scaling only.

        Increase or decrease size of amax history based on given `length`.
659
660
661
662
663
664
665
666
667
668

        .. warning::
            This changes the underlying amax memory location.
        """
        if fwd is None:
            fp8_meta_tensor_keys = ("scaling_fwd", "scaling_bwd")
        else:
            fp8_meta_tensor_keys = ("scaling_fwd" if fwd else "scaling_bwd",)

        for meta_key in fp8_meta_tensor_keys:
669
670
671
            if meta_key not in self.fp8_meta:
                # Handles non-parameter FP8 modules, e.g. DPA.
                continue
672
673
674
675
676
            curr_len = self.fp8_meta[meta_key].amax_history.shape[0]
            if length == curr_len:
                continue
            if length < curr_len:
                self.fp8_meta[meta_key].amax_history = (
677
678
                    self.fp8_meta[meta_key].amax_history[:length].clone()
                )
679
680
681
682
683
684
            elif length > curr_len:
                extra_rows = length - curr_len
                self.fp8_meta[meta_key].amax_history = F.pad(
                    self.fp8_meta[meta_key].amax_history, pad=(0, 0, 0, extra_rows)
                )

685
686
            # Update quantizers with new amax pointers.
            self.quantizers[meta_key] = self.fp8_meta[meta_key].make_quantizers()
687
688
            # Make sure weight tensors has correct quantizers
            self._update_weight_quantizers()
689

690
691
            # Update the global buffers with new amax and history pointers.
            if FP8GlobalStateManager.get_buffer_info() in self.fp8_meta:
692
693
694
                fwd_pos, fwd_key, bwd_pos, bwd_key = self.fp8_meta[
                    FP8GlobalStateManager.get_buffer_info()
                ]
695
696
697
698
699
                for pos, buffer_key in zip((fwd_pos, bwd_pos), (fwd_key, bwd_key)):
                    if buffer_key in FP8GlobalStateManager.global_amax_buffer:
                        assert (
                            buffer_key in FP8GlobalStateManager.global_amax_history_buffer
                        ), "TE internal error during amax history change."
700
701
702
                        FP8GlobalStateManager.global_amax_buffer[buffer_key][pos] = self.fp8_meta[
                            meta_key
                        ].amax_history[0]
703
                        FP8GlobalStateManager.global_amax_history_buffer[buffer_key][pos] = (
704
705
                            self.fp8_meta[meta_key].amax_history
                        )
706

707
    def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
708
709
710
        """Init scales and amaxes for fwd | bwd."""
        fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"

711
        # Return early if recipe state matches recipe
712
        if self.fp8_meta_tensors_initialized:
713
714
715
716
717
718
            recipe_state = self.fp8_meta[fp8_meta_tensor_key]
            if recipe.delayed() and isinstance(recipe_state, DelayedScalingRecipeState):
                self.adjust_amax_history_length(recipe.amax_history_len, fwd=fwd)
                return
            if recipe.mxfp8() and isinstance(recipe_state, MXFP8BlockScalingRecipeState):
                return
719
720
721
722
            if recipe.float8_current_scaling() and isinstance(
                recipe_state, Float8CurrentScalingRecipeState
            ):
                return
723
724
725
726
            if recipe.float8_block_scaling() and isinstance(
                recipe_state, Float8BlockScalingRecipeState
            ):
                return
727
728
            if recipe.nvfp4() and isinstance(recipe_state, NVFP4BlockScalingRecipeState):
                return
729
730
731

        # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
        # 2 (grad_output and grad_input) for bwd
732
        num_fp8_tensors = self.fp8_meta["num_gemms"] * 3 if fwd else self.fp8_meta["num_gemms"] * 2
733

734
735
736
737
738
        # Initialize recipe state and quantizers
        recipe_state = RecipeState.create(
            recipe,
            mode=("forward" if fwd else "backward"),
            num_quantizers=num_fp8_tensors,
739
740
        )

741
742
743
        self.fp8_meta[fp8_meta_tensor_key] = recipe_state
        self.quantizers[fp8_meta_tensor_key] = recipe_state.make_quantizers()

744
745
746
747
748
749
750
751
752
    def _update_weight_quantizers(self) -> None:
        """Update the quantizers for the weight tensors."""
        weight_tensors = self._get_weight_tensors()
        weight_quantizers = self._get_weight_quantizers()
        assert len(weight_tensors) == len(weight_quantizers), (
            f"Number of weight tensors ({len(weight_tensors)}) and quantizers "
            f"({len(weight_quantizers)}) must match"
        )
        for weight, quantizer in zip(weight_tensors, weight_quantizers):
753
            if quantizer is not None and isinstance(weight, QuantizedTensorStorage):
754
755
                weight.update_quantizer(quantizer)

756
    def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage]]:
757
758
759
760
761
762
763
764
765
766
767
        """Get the weight tensors of the module."""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement _get_weight_tensors function"
        )

    def _get_weight_quantizers(self) -> List[Quantizer]:
        """Get the weight quantizers of the module."""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement _get_weight_quantizers function"
        )

768
    def init_fp8_meta_tensors(self, recipe: Recipe) -> None:
769
        """Init scales and amaxes."""
770
771
772
        self.set_meta_tensor(True, recipe)
        self.set_meta_tensor(False, recipe)

773
        self.fast_setattr("fp8_meta_tensors_initialized", True)
774

775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
    def get_fp8_meta_tensors(self) -> None:
        """Get scales and amaxes."""
        fwd_key, bwd_key = "scaling_fwd", "scaling_bwd"
        if fwd_key not in self.fp8_meta or bwd_key not in self.fp8_meta:
            return None

        fp8_meta_tensors = {fwd_key: [], bwd_key: []}
        with torch.no_grad():
            for key in (fwd_key, bwd_key):
                fp8_meta_tensors[key].append(self.fp8_meta[key].scale.clone())
                fp8_meta_tensors[key].append(self.fp8_meta[key].amax_history.clone())
        return fp8_meta_tensors

    def reset_fp8_meta_tensors(self, fp8_meta_tensors=None) -> None:
        """Reset scales and amaxes."""
790

791
792
793
794
795
        def reset(key):
            if key in self.fp8_meta:
                if fp8_meta_tensors is None:
                    self.fp8_meta[key].scale.copy_(torch.ones_like(self.fp8_meta[key].scale))
                    self.fp8_meta[key].amax_history.copy_(
796
797
                        torch.zeros_like(self.fp8_meta[key].amax_history)
                    )
798
799
800
                else:
                    assert key in fp8_meta_tensors, "Cannot reset fp8 tensors."
                    self.fp8_meta[key].scale.copy_(fp8_meta_tensors[key][0])
801
                    self.fp8_meta[key].amax_history.copy_(fp8_meta_tensors[key][1])
802

803
804
805
806
        with torch.no_grad():
            reset("scaling_fwd")
            reset("scaling_bwd")

807
    def get_extra_state(self) -> torch.Tensor:
808
        """Save before checkpointing."""
809

810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
        # This implementation is working around a few issues:
        #
        # (1) PyTorch's "extra state" infrastructure might be able to
        #     support any picklable type, but they make no guarantees.
        #     We have experienced problems (e.g. in ONNX export) with
        #     non-tensor extra state.
        # (2) PyTorch's checkpointing infrastructure does not remap
        #     devices for "extra state" like it does for "state dict".
        #     Thus, we want to avoid putting extra state on the GPU
        #     since it may be loaded on the wrong device.
        # (3) The extra state consists of many small tensors. If we
        #     want to copy them all to CPU, then we need to avoid the
        #     overhead of many GPU-CPU memory transfers.
        #
        # See: https://github.com/NVIDIA/TransformerEngine/pull/351
        # See: https://github.com/NVIDIA/TransformerEngine/pull/363

        def to_cpu(src: torch.Tensor) -> torch.Tensor:
            """Helper function to make CPU copy of tensor

            Memory transfer is asynchronous w.r.t. host, so GPU should
            be synchronized before using result.

            """
            dst = torch.empty_like(src, device="cpu")
            dst.copy_(src, non_blocking=True)
            return dst

        # Store FP8 state if needed
        state = None
840
        fp8_checkpoint = self.fp8_meta["fp8_checkpoint"] or self.fp8 or self.fp8_calibration
841
        if not fp8_checkpoint:
842
            return torch.empty(0, dtype=torch.uint8)
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860

        # Copy tensors to CPU and store
        state = {}
        state["recipe"] = self.fp8_meta["recipe"]
        if state["recipe"].delayed():
            state["scale_fwd"] = to_cpu(self.fp8_meta["scaling_fwd"].scale)
            state["amax_history_fwd"] = to_cpu(self.fp8_meta["scaling_fwd"].amax_history)
            state["scale_bwd"] = to_cpu(self.fp8_meta["scaling_bwd"].scale)
            state["amax_history_bwd"] = to_cpu(self.fp8_meta["scaling_bwd"].amax_history)

        # Store other pickelable values
        extra = {}
        for k, v in self.fp8_meta.items():
            if k != "buffer_index_and_autocast_key" and isinstance(
                v, (bool, int, float, str, tuple, list)
            ):
                extra[k] = v
        state["extra_fp8_variables"] = extra
861

862
863
864
865
        # Serialize state into byte tensor
        torch.cuda.synchronize()
        state_serialized = bytearray(pickle.dumps(state))
        state_serialized = torch.frombuffer(state_serialized, dtype=torch.uint8)
866
        return state_serialized
867

868
    def set_extra_state(self, state: torch.Tensor) -> None:
869
        """Load previous state."""
870
871
872
873
874

        # Maintain backwards compatibility with older checkpoints.
        if state is None:
            return

875
        # Load state
876
        if isinstance(state, torch.Tensor):
877
878
879
            # No FP8 is indicated by an empty tensor we don't need to unpickle.
            if state.numel() == 0:
                return
880
            # Default format: byte tensor with pickled data
881
            state = pickle.loads(state.detach().cpu().numpy().tobytes())
882
        elif isinstance(state, io.BytesIO):
883
            # Deprecated format with io.BytesIO
884
            state.seek(0)
885
            state = torch.load(state, map_location="cuda")
886
887
        else:
            raise RuntimeError("Unsupported checkpoint format.")
888
889
890

        if state is None:
            return
891

892
893
894
895
896
897
898
899
        # TE 1.x checkpoint compatibility: add DelayedScaling recipe if missing
        if "recipe" not in state:
            # TE 1.x only supported delayed scaling, which was the default recipe
            state["recipe"] = DelayedScaling()
            # TE 1.x also saved scale_inv, which is not needed with Recipe object
            state.pop("scale_inv_fwd", None)
            state.pop("scale_inv_bwd", None)

900
        # Load extra items
901
        self.fp8_meta.update(state["extra_fp8_variables"])
902
        self.fp8_meta["recipe"] = state["recipe"]
903
904
905
        if "global_fp8_buffer_pos_fwd_recompute" in self.fp8_meta:
            del self.fp8_meta["global_fp8_buffer_pos_fwd_recompute"]

906
        # Initialize before loading
907
        self.init_fp8_meta_tensors(self.fp8_meta["recipe"])
908
909
910
911
912
913
914
915
916
917
918

        def copy_tensor(src: torch.Tensor, dst: torch.Tensor) -> None:
            """Helper function to copy tensor from CPU

            Memory transfer is asynchronous w.r.t. host, so GPU should
            be synchronized before using result.

            """
            dst.copy_(src, non_blocking=True)

        # Load tensors
919
920
921
922
923
        if self.fp8_meta["recipe"].delayed():
            copy_tensor(state["scale_fwd"], self.fp8_meta["scaling_fwd"].scale)
            copy_tensor(state["amax_history_fwd"], self.fp8_meta["scaling_fwd"].amax_history)
            copy_tensor(state["scale_bwd"], self.fp8_meta["scaling_bwd"].scale)
            copy_tensor(state["amax_history_bwd"], self.fp8_meta["scaling_bwd"].amax_history)
924
        torch.cuda.synchronize()
925
926
927
928
929

    def set_activation_dtype(self, inp: torch.Tensor) -> None:
        """Get activation data type for AMP."""
        # Native AMP (`torch.autocast`) gets highest priority
        if torch.is_autocast_enabled():
930
            self.fast_setattr("activation_dtype", torch_get_autocast_gpu_dtype())
931
932
933
            return

        # All checks after this have already been performed once, thus skip
934
        if self.activation_dtype == inp.dtype:
935
936
            return

937
        dtype = inp.dtype
938
939
940
941
942
943
944
        if not self.allow_different_data_and_param_types:
            for name, param in self.named_parameters():
                if param is not None:
                    assert dtype == param.dtype, (
                        "Data types for parameters must match when outside of autocasted region. "
                        f" Found input dtype: {dtype} and {name!r} dtype: {param.dtype}"
                    )
945
        self.fast_setattr("activation_dtype", dtype)
946
947

    def set_tensor_parallel_group(self, tp_group: Union[dist_group_type, None]) -> None:
948
949
950
951
952
953
        """
        Set the tensor parallel group for the given
        module before executing the forward pass.

        Parameters
        ----------
Paweł Gadziński's avatar
Paweł Gadziński committed
954
        tp_group : ProcessGroup, default = None
955
956
                  tensor parallel process group.
        """
957
958
        self.fast_setattr("tp_group", tp_group)
        self.fast_setattr("tp_group_initialized", True)
959

960
961
962
    def _get_fp8_params(self) -> Union[List[torch.Tensor], None]:
        """returns the FP8 weights."""
        fp8_params = []
963
        for param in self.parameters(recurse=False):
964
            if isinstance(param, QuantizedTensor) and param.requires_grad:
965
966
967
968
969
                fp8_params.append(param)
        if len(fp8_params) == 0:
            return None
        return fp8_params

970
971
    # This routine is shared across FP8 and FP8_calibration paths so should not actually
    # assume FP8 execution.
972
    def init_fp8_metadata(self, num_gemms: int = 1) -> None:
973
        """Initialize fp8 related metadata and tensors during fprop."""
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
        meta = self.fp8_meta

        fp8 = FP8GlobalStateManager.is_fp8_enabled()
        fp8_parameters = FP8GlobalStateManager.with_fp8_parameters()
        fp8_calibration = FP8GlobalStateManager.is_fp8_calibration()
        self.fast_setattr("fp8_parameters", fp8_parameters)
        self.fast_setattr("fp8", fp8)
        self.fast_setattr("fp8_calibration", fp8_calibration)
        fp8_enabled = fp8 or fp8_calibration
        meta["fp8_checkpoint"] = fp8_enabled

        _original_recipe = None

        if fp8_parameters or fp8_enabled:
            _original_recipe = meta.get("recipe", None)
            if self.fp8_initialized and FP8GlobalStateManager.get_fp8_recipe() == _original_recipe:
990
                # FP8 init has already been run and recipe is the same, don't do anything.
991
                return
992
            meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe()
993
994
        else:
            # If fp8 isn't enabled, turn off and return.
995
            self.fast_setattr("fp8_initialized", False)
996
997
            return

998
999
1000
        if fp8_parameters and not self.fp8_initialized:
            meta["num_gemms"] = num_gemms
            self.init_fp8_meta_tensors(meta["recipe"])
1001
1002
1003

        if fp8_enabled:
            # Set FP8 and other FP8 metadata
1004
1005
            meta["num_gemms"] = num_gemms
            meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group()
1006
1007

            # Set FP8_MAX per tensor according to recipe
1008
1009
1010
            if hasattr(meta["recipe"], "fp8_format"):
                meta["fp8_max_fwd"] = meta["recipe"].fp8_format.value.max_fwd
                meta["fp8_max_bwd"] = meta["recipe"].fp8_format.value.max_bwd
1011
1012

            # Allocate scales and amaxes
1013
1014
            self.init_fp8_meta_tensors(meta["recipe"])
            self.fast_setattr("fp8_initialized", True)
1015

1016
            meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe()
1017

1018
        _current_recipe = meta["recipe"]
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
        if _original_recipe is not None and not (
            issubclass(_current_recipe.__class__, _original_recipe.__class__)
            or issubclass(_original_recipe.__class__, _current_recipe.__class__)
        ):
            warnings.warn(
                f"Recipe type changed from {_original_recipe.__class__.__name__} "
                f"to {_current_recipe.__class__.__name__}. "
                "This may affect model behavior."
            )
            # Clear cached workspaces as they were created with the old recipe/quantizer type
            self._fp8_workspaces.clear()

1031
1032
1033
1034
    def prepare_forward(
        self,
        inp: torch.Tensor,
        num_gemms: int = 1,
1035
        allow_non_contiguous: bool = False,
1036
        allow_different_data_and_param_types: bool = False,
1037
1038
1039
1040
1041
1042
    ) -> torch.Tensor:
        """Checks and prepares for FWD execution."""
        self.fast_setattr(
            "allow_different_data_and_param_types", allow_different_data_and_param_types
        )
        self.fast_setattr("forwarded_at_least_once", True)
1043

1044
1045
        # Activation recomputation is used and this is the second forward phase.
        if self.fp8 and in_fp8_activation_recompute_phase():
1046
            delayed_scaling_recipe = self.fp8_meta["recipe"].delayed()
1047
            FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta)
1048
1049
1050
1051
1052
1053
1054
        else:
            assert inp.is_cuda, "TransformerEngine needs CUDA."

            if self.tp_size > 1:
                assert self.tp_group_initialized, "TP group not initialized."

            self.set_activation_dtype(inp)
1055
            self.init_fp8_metadata(num_gemms=num_gemms)
1056
            self._check_weight_tensor_recipe_correspondence()
1057

1058
1059
1060
1061
1062
1063
1064
            delayed_scaling_recipe = self.fp8 and self.fp8_meta["recipe"].delayed()
            if delayed_scaling_recipe:
                if self.sequence_parallel:
                    assert self.fp8_meta["recipe"].reduce_amax, (
                        "Amax reduction across tensor parallel group is "
                        "necessary when using sequence parallelism with FP8."
                    )
1065

1066
1067
                if not FP8GlobalStateManager.fp8_graph_capturing():
                    FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(self.fp8_meta)
1068

1069
1070
1071
                # Activation recomputation is used and this is the first forward phase.
                if self.training and is_fp8_activation_recompute_enabled():
                    FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta)
1072

1073
1074
1075
1076
        nvtx_range_push(self.__class__.__name__ + " forward")
        if not allow_non_contiguous and not inp.is_contiguous():
            inp = inp.contiguous()
        return inp
1077

1078
1079
1080
1081
1082
1083
    def end_forward(self):
        """
        Required to be called at the end of the forward function to properly handle
        DelayedScaling metadata handling and the NVTX ranges.
        """
        delayed_scaling_recipe = self.fp8 and self.fp8_meta["recipe"].delayed()
1084
        if delayed_scaling_recipe and self.fp8 and in_fp8_activation_recompute_phase():
1085
            FP8GlobalStateManager.restore_fp8_meta_tensors(self.fp8_meta)
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
        nvtx_range_pop()

    @contextmanager
    def prepare_forward_ctx(
        self,
        inp: torch.Tensor,
        num_gemms: int = 1,
        allow_non_contiguous: bool = False,
        allow_different_data_and_param_types: bool = False,
    ) -> Generator[torch.Tensor, None, None]:
        """Checks and prepares for FWD execution."""
        inp = self.prepare_forward(
            inp, num_gemms, allow_non_contiguous, allow_different_data_and_param_types
        )
        try:
            yield inp
        finally:
            self.end_forward()
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123

    def set_nccl_overlap_warning_if_tp(self) -> None:
        """When using TP, the NCCL communication needs to be scheduled
        before the GEMM for there to be a guaranteed overlap. From the
        host side in TE, the comm calls are always launched first, but
        to ensure that the GEMM isn't scheduled first, the environment
        variable `CUDA_DEVICE_MAX_CONNECTIONS` needs to be set to 1 to
        force a single channel.
        """
        if self.tp_size == 1:
            return
        num_cuda_work_queues = int(os.getenv("CUDA_DEVICE_MAX_CONNECTIONS", "0"))
        if num_cuda_work_queues != 1:
            warnings.warn(
                "To guarantee overlapping TP and SP collectives with the backward"
                "GEMMs, set environment variable CUDA_DEVICE_MAX_CONNECTIONS = 1"
            )

    @staticmethod
    def grad_output_preprocess(
1124
1125
1126
1127
        ctx,
        grad_output: torch.Tensor,
        row_parallel_mode: bool,
        quantizer: Optional[Quantizer],
1128
1129
1130
    ) -> Tuple[Union[torch.Tensor, None], ...]:
        """Utility function for backward.
        Returns tuple in order (all optional/None based on training precion/recipe):
1131
1132
            R1: gathered `grad_output`.
            R2: bias gradient on R1.
1133
1134

        """
1135
1136
        grad_output = grad_output.reshape((-1, grad_output.shape[-1]))
        grad_output = grad_output.contiguous()
1137
1138
        gather_grad_output = row_parallel_mode and ctx.sequence_parallel

1139
        # Non-FP8 case: bgrad is fused with wgrad for this case.
1140
        if not ctx.fp8 and not ctx.debug:
1141
            if gather_grad_output:
1142
                if not ctx.ub_overlap_ag:  # Perform NCCL all-gather
1143
                    grad_output, _ = gather_along_first_dim(grad_output, ctx.tp_group)
1144
1145
1146
1147
1148
1149
1150
                else:  # Initialize Userbuffers all-gather
                    grad_output, _ = fill_userbuffers_buffer_for_all_gather(
                        ctx.ub_obj_gradout,
                        grad_output,
                        None,
                        ctx.tp_group,
                    )
1151
1152
1153
            return grad_output, None

        # FP8 with all-gather: unfused bgrad, fused cast + transpose
1154
        # Also supports debug quantization, which is handled inside gather_along_first_dim.
1155
1156
        if gather_grad_output:
            grad_bias = None
1157
            if ctx.use_bias:
1158
                grad_bias = grad_output.view(-1, grad_output.shape[-1]).sum(dim=0)
1159
            if ctx.ub_overlap_ag:
1160
1161
                # Quantize the gradient if needed
                if not isinstance(
1162
1163
1164
                    grad_output,
                    (
                        QuantizedTensor,
1165
1166
1167
                        Float8TensorStorage,
                        MXFP8TensorStorage,
                        Float8BlockwiseQTensorStorage,
1168
                    ),
1169
1170
1171
1172
                ):
                    grad_output = quantizer(grad_output)

                # Copy into communication buffer, and replace original gradient with it
1173
1174
1175
1176
1177
1178
                grad_output, _ = fill_userbuffers_buffer_for_all_gather(
                    ctx.ub_obj_gradout,
                    grad_output,
                    quantizer,
                    ctx.tp_group,
                )
1179
            else:
1180
1181
1182
1183
                grad_output, _ = gather_along_first_dim(
                    grad_output,
                    ctx.tp_group,
                    quantizer=quantizer,
1184
                )
1185
            return grad_output, grad_bias
1186

1187
1188
1189
1190
        # Debug without all-gather: unfused cast and bgrad
        # bgrad only if wgrad is in FP8, otherwise it is fused with wgrad and we return None
        if ctx.debug:
            grad_output_ = quantizer(grad_output)
1191
            if ctx.use_bias:
1192
1193
1194
1195
1196
1197
                grad_bias = grad_output.view(-1, grad_output.shape[-1]).sum(dim=0)
            else:
                grad_bias = None
            grad_output = grad_output_
            return grad_output, grad_bias

1198
1199
        # FP8 without all-gather: fused bgrad + cast + transpose
        grad_bias = None
1200
        if ctx.use_bias:
1201
1202
            if isinstance(
                grad_output,
1203
1204
1205
1206
1207
1208
                (
                    QuantizedTensor,
                    Float8TensorStorage,
                    MXFP8TensorStorage,
                    Float8BlockwiseQTensorStorage,
                ),
1209
            ):
1210
                grad_bias = grad_output.dequantize().view(-1, grad_output.shape[-1]).sum(dim=0)
1211
            else:
1212
                if isinstance(quantizer, Float8BlockQuantizer):
1213
1214
1215
1216
                    # unfuse bgrad for now until cast_transpose + dgrad calculation is ready for Float8BlockQuantizer.
                    grad_bias = grad_output.view(-1, grad_output.shape[-1]).sum(dim=0)
                else:
                    grad_bias, grad_output = tex.bgrad_quantize(grad_output, quantizer)
1217
        if not isinstance(grad_output, QuantizedTensorStorage):
1218
1219
            grad_output = quantizer(grad_output)
        return grad_output, grad_bias
1220

1221
1222
1223
1224
1225
1226
    def register_parameter(self, name, param, **kwargs):
        """
        Thin wrapper around PyTorch parameter registration to stash additional parameter
        metedata used in deferred initialization.
        """
        super().register_parameter(name, param)
1227
1228
1229
1230
1231
1232
        # Initialize param_init_meta exactly once during the init. FSDP2 can call
        # register parameter again to change parameters to DTensors. And it calls
        # it without custom fp8 specific kwargs that we need. And so we dont want
        # to reset/loose our fp8 init attributes.
        if hasattr(self, "param_init_meta") and name not in self.param_init_meta:
            self.param_init_meta[name] = _ParameterInitMeta(**kwargs)
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243

    def reset_parameters(self, defer_init: Optional[bool] = False) -> None:
        """
        Reset all module parameters to initial values. Unless deferred initialization
        is specified, all parameters on a 'meta' device are also materialized on a real cuda
        device before the values are reset to initial.
        """
        if defer_init:
            return

        for name, param in self.named_parameters(recurse=False):
1244
1245
1246
1247
1248
            # Check if parameter is a DTensor (FSDP2) or regular tensor
            is_dtensor = isinstance(param, DTensor)
            dtensor_param = param if is_dtensor else None
            # Need to update/quantize local tensor in case of DTensor
            param = param._local_tensor if is_dtensor else param
1249
            # Ensure parameter is on a real device
1250
1251
            if param.device == torch.device("meta"):
                param = torch.empty_like(param, device="cuda")
1252
1253
1254
1255
1256
1257
            # Initialize the parameter values on device
            init_fn = self.param_init_meta[name].init_fn
            get_rng_state_tracker = self.param_init_meta[name].get_rng_state_tracker
            if get_rng_state_tracker is None:
                init_fn(param)
            else:
1258
1259
1260
1261
1262
1263
                if hasattr(self, "rng_tracker_name") and self.rng_tracker_name:
                    with get_rng_state_tracker().fork(self.rng_tracker_name):
                        init_fn(param)
                else:
                    with get_rng_state_tracker().fork():
                        init_fn(param)
1264

1265
            # Wrap parameters in QuantizedTensor if needed
1266
            fp8_meta_index = self.param_init_meta[name].fp8_meta_index
1267
            high_precision_init_val = None
1268
            if self.primary_weights_in_fp8 and fp8_meta_index is not None:
1269
1270

                # Keep high-precision values on CPU if needed
1271
1272
1273
                if self.preserve_high_precision_init_val:
                    high_precision_init_val = param.detach().cpu()

1274
                # Configure quantizer
1275
                quantizer = self.quantizers["scaling_fwd"][fp8_meta_index]
1276
1277
1278
                if quantizer is None:
                    raise RuntimeError("Weight quantizer has not been initialized")
                quantizer.set_usage(rowwise=True, columnwise=torch.is_grad_enabled())
1279
                quantizer.internal = False
1280
1281
1282
1283
1284
1285
1286
1287
1288
                if is_dtensor and isinstance(quantizer, Float8CurrentScalingQuantizer):
                    device_mesh = dtensor_param.device_mesh
                    amax_reduction_group = (
                        device_mesh.get_group(mesh_dim="shard")
                        if device_mesh.ndim > 1
                        else device_mesh.get_group()
                    )
                    quantizer.amax_reduction_group = amax_reduction_group
                    quantizer.with_amax_reduction = True
1289
                # Quantize parameter
1290
                param = quantizer(param)
1291
1292
1293
1294
1295

            # Redo parameter wrap in case we broke it above
            # NOTE: Currently this can only be broken when primary weights are in Fp8 but
            #       re-applying the nn.Parameter() wrap is a no-op when the input is already
            #       a parameter so we always re-apply it just for extra safety.
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
            if is_dtensor:
                # recreate the DTensor from the parameter.
                dtensor_param = DTensor.from_local(
                    param,
                    device_mesh=dtensor_param.device_mesh,
                    placements=dtensor_param.placements,
                    shape=dtensor_param.size(),
                    stride=dtensor_param.stride(),
                )
                dtensor_param = torch.nn.Parameter(dtensor_param)
            else:
                param = torch.nn.Parameter(param)
1308
1309

            # Keep high-precision values on CPU if needed
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
            if high_precision_init_val is not None:

                # - Master weights are initialized from model weights, if we use fp8 primary
                #   weights to initialize master weights, the numerical values of master weights
                #   are not consistent with the numerical values when we initialize them from
                #   bf16/fp16 weights.
                # - So we add a `_high_precision_init_val` attribute to each model weight to store
                #   the original bf16/fp16 weight on cpu before casting it to fp8. And users can
                #   use `get_high_precision_init_val` to get this cpu tensor.
                # - This cpu tensor is not needed once the master weight is initialized, so users
                #   should call `clear_high_precision_init_val` to remove it after master weight
                #   is initialized.

                def get(self):
                    if hasattr(self, "_high_precision_init_val"):
                        return self._high_precision_init_val
                    return None

                def clear(self):
                    if hasattr(self, "_high_precision_init_val"):
                        del self._high_precision_init_val

                param._high_precision_init_val = high_precision_init_val
                param.get_high_precision_init_val = MethodType(get, param)
                param.clear_high_precision_init_val = MethodType(clear, param)
1335
                # Update the parameter based on its type
1336

1337
            if not is_dtensor:
1338
                self.module_setattr(name, param)
1339
            else:
1340
                self.module_setattr(name, dtensor_param)
1341

1342
1343
1344
    @abstractmethod
    def forward(self):
        """Needs override."""
1345

1346
    def get_weight_workspace(
1347
        self,
1348
1349
        *,
        tensor: Optional[torch.Tensor] = None,
1350
        quantizer: Optional[Quantizer] = None,
1351
1352
1353
        cache_name: Optional[str] = None,
        update_workspace: bool = True,
        skip_update_flag: Optional[torch.Tensor] = None,
1354
        fsdp_group: Optional[dist_group_type] = None,
1355
        workspace_dtype: Optional[torch.dtype] = None,
1356
    ) -> QuantizedTensor:
1357
        """Get workspace buffer for weights and maybe update its values
1358
1359
1360
1361
1362
1363
1364
1365

        The workspace buffer may be cached for future function calls.

        Parameters
        ----------
        tensor : torch.Tensor, optional
            Values to copy into workspace. Required if the workspace
            is being constructed or updated.
1366
1367
1368
        quantizer: Quantizer, optional
            Quantizer used to cast the weights. Required if the
            workspace is being constructed or updated.
1369
1370
        cache_name: str, optional
            Key for caching.
Paweł Gadziński's avatar
Paweł Gadziński committed
1371
        update_workspace: bool, default = True
1372
1373
1374
1375
            Update workspace with values from `tensor`.
        skip_update_flag: torch.Tensor, optional
            GPU flag to skip updating the workspace. Take precedence
            over `update_workspace` if provided.
1376
1377
        fsdp_group: bool, default = None
            FSDP process group that the weights are distributed over.
1378
1379
1380
        workspace_dtype: torch.dtype, default = None
            If weight workspace contains high-precision tensor - for example
            for debug quantization, this is dtype of the tensor.
1381
1382
        """

1383
1384
1385
        # Handle case where weights are already quantized
        # Note: Make sure weights have required usages, but do not
        # destroy unnecessary usages since they may be used later.
1386
        if isinstance(tensor, QuantizedTensor):
1387
1388
1389
1390
1391
1392
            update_rowwise_usage = True if quantizer.rowwise_usage else None
            update_columnwise_usage = True if quantizer.columnwise_usage else None
            tensor.update_usage(
                rowwise_usage=update_rowwise_usage,
                columnwise_usage=update_columnwise_usage,
            )
1393
1394
            return tensor

1395
        # Try getting workspace from cache
1396
1397
1398
        out = None
        if cache_name is not None:
            out = self._fp8_workspaces.get(cache_name, None)
1399
1400
1401
1402

        # Reset cache if workspace is invalid
        if out is not None and quantizer is not None:
            reset_cache = False
1403
            if isinstance(out, Float8TensorStorage):
1404
1405
1406
1407
1408
1409
                if (
                    not is_non_tn_fp8_gemm_supported()
                    and quantizer.columnwise_usage
                    and out._transpose is None
                ):
                    reset_cache = True
1410
            elif isinstance(out, MXFP8TensorStorage):
1411
                if quantizer.rowwise_usage and out._rowwise_data is None:
1412
                    reset_cache = True
1413
                elif quantizer.columnwise_usage and out._columnwise_data is None:
1414
                    reset_cache = True
1415
1416
1417
1418
1419
            elif isinstance(out, NVFP4TensorStorage):
                if quantizer.rowwise_usage and out._rowwise_data is None:
                    reset_cache = True
                elif quantizer.columnwise_usage and out._columnwise_data is None:
                    reset_cache = True
1420
1421
1422
            if isinstance(out, DebugQuantizedTensor) != isinstance(quantizer, DebugQuantizer):
                reset_cache = True
            if reset_cache:
1423
                out = None
1424
                del self._fp8_workspaces[cache_name]
1425

1426
1427
1428
1429
1430
        # Gather cached Fp8 workspace if it's distributed
        # NOTE: FSDP sharding is supported only for Fp8 buffers and will not work
        #       for models initialized with Fp8 primary weights.
        if (
            out is not None
1431
            and tensor is not None
1432
            and fsdp_group is not None
1433
            and out.data.shape != tensor.data.shape
1434
1435
1436
1437
        ):
            _fsdp_gather_tensors(fsdp_group, [tensor.data.shape], out)

        # Construct workspace if needed
1438
        if out is None:
1439
            if tensor is None or quantizer is None:
1440
                raise ValueError(
1441
                    "tensor and quantizer kwargs must be provided to construct FP8 workspace"
1442
                )
1443
1444
1445
1446
1447
1448
1449

            if cache_name is not None:
                # Ensure the tensor in the cache is an instance of torch.Tensor,
                # as it persists beyond a single forward pass.
                # Setting internal=True would cause the data to be removed in prepare_for_saving(...).
                quantizer_internal = quantizer.internal
                quantizer.internal = False
1450
            out = quantizer.quantize(tensor, dtype=workspace_dtype)
1451
1452
            if cache_name is not None:
                quantizer.internal = quantizer_internal
1453
1454

            # Update cache
1455
1456
            if cache_name is not None:
                self._fp8_workspaces[cache_name] = out
1457
            return out
1458
1459
1460
1461
1462
1463

        # Update workspace if needed
        if skip_update_flag is not None:
            update_workspace = True
        if update_workspace:
            if tensor is None:
1464
                raise ValueError("tensor kwarg must be provided to update FP8 workspace")
1465
            if hasattr(out, "quantize_"):
1466
                out.quantize_(tensor, noop_flag=skip_update_flag)
1467
1468
            else:
                tex.quantize(tensor, quantizer, out, skip_update_flag)
1469
        return out
1470

1471
1472
1473
    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
        """
        This function loads tensors and extra state including fp8 metadata.
        This metadata is essential for copying fp8 tensors, as the copy_ function
        uses the scale_inv parameter from fp8_meta to set the correct scaling factor
        for the new tensor.
        Hence, this extra state must be loaded before the tensor copying process,
        not after, as is typically done in _load_from_state_dict.
        Tensors are copied into fp8 tensors only when self.primary_weights_in_fp8=True,
        otherwise, this behavior is not required.
        """
        if self.primary_weights_in_fp8:
            extra_state_key = prefix + torch.nn.modules.module._EXTRA_STATE_KEY_SUFFIX
            if extra_state_key in state_dict:
                self.set_extra_state(state_dict[extra_state_key])
1488
1489
1490
        super()._load_from_state_dict(
            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
        )
1491

1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
    def register_wgrad_accumulation_and_reduce_hooks(self, wgrad_accumulation_and_reduce_hook):
        """
        This method is used to manually control the weight gradient accumulation and reduce.
        This method should be called before the backward() method.
        Set the skip_wgrad_accumulation_and_reduce to True to skip the weight gradient accumulation
        and reduce in backward();
        And register the wgrad_accumulation_and_reduce_func to be called in backward_dw() method.
        """
        self.wgrad_accumulation_and_reduce_hooks.append(wgrad_accumulation_and_reduce_hook)

1502
1503
1504
1505
1506
1507
1508
1509
1510
    def need_backward_dw(self):
        """
        Check if this module needs to execute the delayed weight gradient computation.
        This method should be used at the beginning of self.backward_dw() to determine if it
        should actually be executed or just return without doing anything.
        User can also manually call this method to check that before calling into backward_dw().
        """
        return self.wgrad_store is not None and self.wgrad_store.delay_wgrad_compute()

1511
1512
1513
1514
1515
    def backward_dw(self):
        """
        Execute the delayed weight gradient computation.
        This method is called after the main backward pass to compute weight gradients.
        """
1516
        if not self.need_backward_dw():
1517
            return
1518
        with get_nvtx_range_context(f"_{self.__class__.__name__}_wgrad"):
1519
            (wgrad, bgrad), _ = self.wgrad_store.pop()
1520
            if not self.fuse_wgrad_accumulation:
1521
                weight_tensor = noop_cat(self._get_weight_tensors())
1522
                weight_tensor.grad = wgrad.to(weight_tensor.dtype)
1523
1524
1525
            if self.use_bias:
                bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names])
                if bias_tensor.grad is None:
1526
                    bias_tensor.grad = bgrad.to(bias_tensor.dtype)
1527
1528
1529
1530
            del wgrad
            del bgrad
            for wgrad_accumulation_and_reduce_hook in self.wgrad_accumulation_and_reduce_hooks:
                wgrad_accumulation_and_reduce_hook()
1531

1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
    def is_debug_iter(self) -> bool:
        """
        This function checks if the debug should be enabled for this layer.
        """
        debug = TEDebugState.debug_enabled
        if not debug:
            return False

        # If layer is run first time in new iteration,
        # we need to check if the debug should be enabled for this layer -
        # maybe in previous iterations debug features returned information
        # that no feature will be active for this layer for multiple next iterations.
        started_new_iteration = TEDebugState.get_iteration() != getattr(
            self, "debug_last_iteration", None
        )
        if started_new_iteration:
            if self.next_iter_when_debug_should_be_run is None:
                debug = False
            else:
                debug = TEDebugState.get_iteration() >= self.next_iter_when_debug_should_be_run
1552
1553
            self.fast_setattr("debug_last_iteration", TEDebugState.get_iteration())
            self.fast_setattr("debug_enabled_in_this_iteration", debug)
1554
1555
1556
1557
1558
        else:
            # If this is the same iteration as previous invocation of the module,
            # we use the debug value from the first invocation in the iteration.
            debug = self.debug_enabled_in_this_iteration

1559
        self.fast_setattr("debug_last_iteration", TEDebugState.get_iteration())
1560
1561
1562
1563
1564

        if self.wgrad_store is not None:
            if debug and self.wgrad_store.delay_wgrad_compute():
                raise RuntimeError("Delayed wgrad compute is not supported in debug mode.")

1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
        return debug

    def no_debug_features_active(self, quantizers):
        """
        Checks if any debug feature is active for this layer.
        """
        run_current = any_feature_enabled(quantizers)

        # Sometimes features inform that they will not be enabled for particular layer
        # for multiple next iterations.
1575
1576
1577
        self.fast_setattr(
            "next_iter_when_debug_should_be_run", next_iter_when_debug_should_be_run(quantizers)
        )
1578
1579
1580
1581
1582
1583
1584
1585

        if not run_current:
            return True

        if self.primary_weights_in_fp8:
            raise RuntimeError("FP8 weights are not supported in debug mode.")
        return False

1586
1587
1588
    def _validate_name(self):
        """
        Validate name passed to the module.
1589
1590
        It creates a default name with layer count as the variable
        which may be changed by the user of the module.
1591
        """
1592
1593
        if self.name is not None:
            return
1594
1595

        self.name = f"Layer_{TEDebugState.get_layer_count()}"
1596

1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
    def _check_weight_tensor_recipe_correspondence(self) -> None:
        """
        Verify that the weight tensor types match their corresponding recipe type.
        This is invoked in the forward().

        This establishes a 1:1 correspondence between recipe types and tensor types:
        - DelayedScaling → Float8Tensor
        - Float8CurrentScaling → Float8Tensor
        - MXFP8BlockScaling → MXFP8Tensor
        - Float8BlockScaling → Float8BlockTensor

1608
1609
        Example case to check: recipe is DelayedScaling (DelayedScaling is set in autocast()),
        but the weight tensor is MXFP8Tensor (MXFP8BlockScaling is set in quantized_model_init()).
1610
1611
1612
        """
        if not self.fp8 and not self.fp8_calibration:
            return
1613
1614
        if not self.primary_weights_in_fp8:
            return
1615
1616
1617
1618
1619
1620
        if not hasattr(self, "weight_names") or not self.weight_names:
            return

        recipe = self.fp8_meta["recipe"]
        weight_tensors = [getattr(self, name) for name in self.weight_names]
        for i, tensor in enumerate(weight_tensors):
1621
            if isinstance(tensor, QuantizedTensorStorage):
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
                quantizer = tensor._get_quantizer()
                if quantizer is None:
                    continue
                compatible_recipe_class = quantizer._get_compatible_recipe()
                if compatible_recipe_class is None:
                    continue
                if not isinstance(recipe, compatible_recipe_class):
                    raise RuntimeError(
                        f"Recipe mismatch for '{self.weight_names[i]}': tensor supports recipe"
                        f" {compatible_recipe_class.__name__}, but got {recipe.__class__.__name__}."
1632
1633
                        " Please check the recipes assigned during quantized_model_init() and"
                        " autocast() calls."
1634
                    )