base.py 51.9 KB
Newer Older
1
# Copyright (c) 2022-2025, 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
8
9
10
import os
import pickle
import warnings
from abc import ABC, abstractmethod
11
from typing import Any, Dict, Generator, List, Optional, Set, Tuple, Union
12
from contextlib import contextmanager
13
import logging
14
from types import MethodType
15
16
17
18

import torch
import torch.nn.functional as F

19
import transformer_engine_torch as tex
20
21
from transformer_engine.common.recipe import Recipe

22
from ._common import _ParameterInitMeta, noop_cat
23
from ..fp8 import (
24
25
    MXFP8BlockScalingRecipeState,
    DelayedScalingRecipeState,
26
    Float8CurrentScalingRecipeState,
27
    Float8BlockScalingRecipeState,
28
    FP8GlobalStateManager,
29
    RecipeState,
30
31
32
33
34
)
from ..distributed import (
    gather_along_first_dim,
    is_fp8_activation_recompute_enabled,
    in_fp8_activation_recompute_phase,
35
    _fsdp_gather_tensors,
36
37
)
from ..constants import dist_group_type
38
from ..tensor import QuantizedTensor, Quantizer
39
from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer
40
41
from ..tensor._internal.float8_tensor_base import Float8TensorBase
from ..tensor._internal.mxfp8_tensor_base import MXFP8TensorBase
42
from ..tensor._internal.float8_blockwise_tensor_base import Float8BlockwiseQTensorBase
43
44
45
from ...common.recipe import Recipe
from ...debug.pytorch.debug_state import TEDebugState
from ...debug.pytorch.debug_quantization import DebugQuantizer, DebugQuantizedTensor
46

47
48
__all__ = ["initialize_ub", "destroy_ub"]

49
50
51
_2X_ACC_FPROP = False
_2X_ACC_DGRAD = True
_2X_ACC_WGRAD = True
52
_multi_stream_cublas_workspace = []
53
_dummy_wgrads = {}
54
55
56
_cublas_workspace = None
_ub_communicators = None
_NUM_MAX_UB_STREAMS = 3
57
_MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY = None, None
58
layers_atomic_ring_exchange = []
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77


def get_cublas_workspace_size_bytes() -> None:
    """Return 32 MiB if using hopper, 4 MiB for all other architectures."""
    if torch.cuda.get_device_properties(torch.cuda.current_device()).major >= 9:
        return 33_554_432
    return 4_194_304


def get_workspace() -> torch.Tensor:
    """Returns workspace for cublas."""
    global _cublas_workspace
    if _cublas_workspace is None:
        _cublas_workspace = torch.empty(
            get_cublas_workspace_size_bytes(), dtype=torch.uint8, device="cuda"
        )
    return _cublas_workspace


78
79
80
81
def get_multi_stream_cublas_workspace() -> List[torch.Tensor]:
    """Returns workspace for multi-stream cublas."""
    global _multi_stream_cublas_workspace
    if not _multi_stream_cublas_workspace:
82
        for _ in range(tex._num_cublas_streams):
83
84
85
86
87
88
            _multi_stream_cublas_workspace.append(
                torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device="cuda")
            )
    return _multi_stream_cublas_workspace


89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
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()


105
106
def initialize_ub(
    shape: list,
107
    tp_size: int,
108
    use_fp8: bool = False,
109
    dtype: torch.dtype = torch.bfloat16,
110
    ub_cfgs: Optional[dict] = None,
111
    bootstrap_backend: Union[str, torch.distributed.Backend] = None,
112
) -> None:
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
    r"""
    Initialize the Userbuffers communicator for overlapping tensor-parallel communications with
    GEMM compute in te.Linear, te.LayerNormLinear and te.LayerNormMLP modules.

    Parameters
    ----------
    shape : list
            shape of the communication buffer, typically set to be the same as the global shape of
            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)`
    tp_size : int
              number of GPUs in the tensor-parallel process group
    use_fp8 : bool = False
              allocate the communication buffer for FP8 GEMM inputs/outputs
    dtype : torch.dtype = torch.bfloat16
            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",
             "proj_fprop", "proj_dgrad", "proj_wgrad", "fc1_fprop", "fc1_dgrad", "fc2_dgrad",
             "fc2_fprop", "fc2_dgrad"]`.
    bootstrap_backend : str = None
                        `torch.distributed` communication backend for the all-gather, broadcast and
                        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
                        not available. Setting `NVTE_UB_WITH_MPI=1` when building TE overrides this
                        option and always initializes Userbuffers with direct MPI calls in C++,
                        which also requires `MPI_HOME=/path/to/mpi/root` to be set at compile time.
    """
160
    if not tex.device_supports_multicast():
161
        assert bool(int(os.getenv("UB_SKIPMC", "0"))), (
162
163
164
165
            "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."
        )

166
167
168
    global _ub_communicators
    assert _ub_communicators is None, "UB communicators are already initialized."
    _ub_communicators = {}
169
170

    if tex.ubuf_built_with_mpi():
171
172
        # 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...
173
        assert torch.distributed.is_mpi_available()
174
175
        _ = torch.distributed.new_group(backend="mpi")
        helper = tex.CommOverlapHelper()
176
    else:
177
178
        # Bootstrapping with torch.distributed API, so check backend and construct
        # intra/inter-node process groups...
179
180
181
182
183
        assert (
            torch.distributed.is_initialized()
        ), "torch.distributed must be initialized before Userbuffers"
        if bootstrap_backend is None:
            bootstrap_backend = "nccl"
184
            if torch.distributed.is_mpi_available():
185
                bootstrap_backend = "mpi"
186
187
            elif torch.distributed.is_gloo_available():
                bootstrap_backend = "gloo"
188
        else:
189
190
191
192
193
194
195
196
197
            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."
            )
198
199
200
201
202

        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)

203
204
        num_domains = world_size // tp_size
        mydomain_idx = world_rank // tp_size
205
        if num_domains > 1:
206
207
208
209
            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(
210
211
                ranks_per_domain_list, backend=bootstrap_backend
            )
212
213
            local_rank = torch.distributed.get_rank(tp_domain_group)
            tp_domain_ranks = torch.distributed.get_process_group_ranks(tp_domain_group)
214

215
            helper = tex.CommOverlapHelper(world_group, tp_domain_group)
216
        else:
217
218
            # TP model on single NVLink domain, no replication, no data-parallelism
            mydomain_idx = 0
219
            local_rank = world_rank
220
            tp_domain_ranks = list(range(world_size))
221
222

            helper = tex.CommOverlapHelper(world_group)
223

224
        if world_rank == 0:
225
            print(f"!!! [UB] Number of TP domains: {num_domains}\n", end="", flush=True)
226
227
        if local_rank == 0:
            print(
228
                f"!!! [UB] Global ranks on TP domain {mydomain_idx}: {tp_domain_ranks}\n",
229
230
231
232
                end="",
                flush=True,
            )

233
    # Allocate cuBLAS workspace with expanded size for chunking in overlapping GEMM calls
234
    global _cublas_workspace
235
236
237
238
239
240
241
    if _cublas_workspace is None:
        _cublas_workspace = get_workspace().repeat(_NUM_MAX_UB_STREAMS)
    elif _cublas_workspace.numel() != get_cublas_workspace_size_bytes() * _NUM_MAX_UB_STREAMS:
        # This ensures we don't do `.repeat()` on an already expanded workspace
        _cublas_workspace = torch.empty(
            get_cublas_workspace_size_bytes(), dtype=torch.uint8, device="cuda"
        ).repeat(_NUM_MAX_UB_STREAMS)
242
243

    # Default buffer precision: AllGather buffers use fp8 when using fp8 recipe
244
    layers_all_gather_overlap = [
245
246
247
248
249
250
        "qkv_fprop",
        "qkv_dgrad",
        "proj_dgrad",
        "fc1_fprop",
        "fc1_dgrad",
        "fc2_dgrad",
251
    ]
252
    layers_reduce_scatter_overlap = ["proj_fprop", "fc2_fprop", "qkv_wgrad", "fc1_wgrad"]
Jaemin Choi's avatar
Jaemin Choi committed
253
    dgrad_reduce_scatter_overlap = ["qkv_dgrad", "fc1_dgrad"]
254
255
    # Default overlap methods for layers
    methods = {
256
257
258
        "ring_exchange": ["qkv_fprop", "fc1_fprop", "proj_dgrad", "fc2_dgrad"],
        "pipeline": ["proj_fprop", "fc2_fprop"],
        "bulk": ["qkv_dgrad", "qkv_wgrad", "fc1_dgrad", "fc1_wgrad"],
259
260
    }

261
    # AG-RS overlap pairs of layers forming a tensor-parallel block
262
263
    ag_rs_pairs = {"qkv_fprop": "proj_fprop", "fc1_fprop": "fc2_fprop"}
    rs_ag_pairs = {v: k for k, v in ag_rs_pairs.items()}
264
265
266
    global layers_atomic_ring_exchange
    layers_atomic_ring_exchange = []

267
268
269
270
271
272
    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.")

273
    def get_default_config(name):
274
        global _MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY
275
276
        method = get_method(name)
        is_reduce_scatter = name in layers_reduce_scatter_overlap
277
278
        if _MIN_STREAM_PRIORITY is None or _MAX_STREAM_PRIORITY is None:
            _MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY = tex.get_stream_priority_range()
279
280
281
282
283
        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,
284
285
            "set_sm_margin": not method == "ring_exchange",
            "num_splits": tp_size if method == "ring_exchange" else 4,
286
287
288
289
            "aggregate": False,
            "atomic_gemm": False,
            "use_ce": True,
            "fp8_buf": name in layers_all_gather_overlap,
290
291
292
            "comm_priority": _MAX_STREAM_PRIORITY,
            "gemm_priority": _MIN_STREAM_PRIORITY,
            "pipeline_rs_overlap_first_gemm": False,
293
294
295
        }
        return default_cfg

296
297
298
    def add_ub(
        name: str,
        method: str,
299
        is_reduce_scatter: bool,
300
301
        num_sm: int = 16,
        cga_size: int = 2,
302
        set_sm_margin: bool = False,
303
        num_splits: int = 0,
304
305
        aggregate: bool = False,
        atomic_gemm: bool = False,
306
        use_ce: bool = True,
307
        fp8_buf: bool = False,
308
309
310
        comm_priority: int = 0,
        gemm_priority: int = 0,
        pipeline_rs_overlap_first_gemm: bool = False,
311
    ) -> None:
312
313
314
315
316
        if atomic_gemm:
            warnings.warn(
                "Atomic GEMM uses a beta API from cublas and is not tested for all use cases."
            )
            assert use_fp8, "Atomic GEMM overlap supported only for FP8 GEMM."
317
            if method == "bulk":
318
                warnings.warn(
319
                    f"At {name}, atoimic GEMM not is supported for a bulk overlap."
320
321
322
                    "Defaulting to `atomic_gemm=False`."
                )
                atomic_gemm = 0
323
        if not is_reduce_scatter and method == "pipeline":
324
            raise ValueError(
325
                f"At {name}, `pipeline` overlap method is not supported for AllGather."
326
            )
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
        # 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

346
        buffer_dtype = torch.uint8 if (use_fp8 and fp8_buf) else dtype
347
        if method == "ring_exchange":
348
349
350
351
            ub_obj = tex.CommOverlapP2P(
                shape,  # Communication buffer shape
                buffer_dtype,  # Communication buffer data type
                helper,  # Helper for torch.distributed callbacks during bootstrapping
352
                tp_size,  # Tensor-parallel group size (may be different than local_size)
353
354
355
356
357
358
359
360
                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,
361
362
                gemm_priority=gemm_priority,
                comm_priority=comm_priority,
363
            )
364
        else:
365
366
367
368
            ub_obj = tex.CommOverlap(
                shape,  # Communication buffer shape
                buffer_dtype,  # Communication buffer data type
                helper,  # Helper for torch.distributed callbacks during bootstrapping
369
                tp_size,  # Tensor-parallel group size (may be different than local_size)
370
371
372
373
374
375
                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,
376
377
378
                gemm_priority=gemm_priority,
                comm_priority=comm_priority,
                rs_overlap_first_gemm=pipeline_rs_overlap_first_gemm,
379
            )
380
381
        _ub_communicators[name] = ub_obj

Jaemin Choi's avatar
Jaemin Choi committed
382
383
    if ub_cfgs is not None:
        for name in dgrad_reduce_scatter_overlap:
384
385
            if name in ub_cfgs and "method" in ub_cfgs[name] and ub_cfgs[name]["method"] != "bulk":
                wgrad_name = name.replace("dgrad", "wgrad")
Jaemin Choi's avatar
Jaemin Choi committed
386
387
                assert wgrad_name not in ub_cfgs
                layers_reduce_scatter_overlap.remove(wgrad_name)
388
                layers_all_gather_overlap.remove(name)
Jaemin Choi's avatar
Jaemin Choi committed
389
                layers_reduce_scatter_overlap.append(name)
390
391
392
                methods["bulk"].remove(name)
                new_method = ub_cfgs[name]["method"]
                methods[new_method].append(name)
Jaemin Choi's avatar
Jaemin Choi committed
393

394
    for name in methods["ring_exchange"] + methods["pipeline"] + methods["bulk"]:
395
        ub_cfg = get_default_config(name)
396
        if ub_cfgs is not None and name in ub_cfgs:
397
            fp8_buf = (name in layers_all_gather_overlap) or (
398
                ub_cfgs[name].get("fp8_buf", False) and name in methods["pipeline"]
399
            )
400
401
402
            ub_cfg.update(ub_cfgs[name])
            ub_cfg["fp8_buf"] = fp8_buf
        add_ub(name, **ub_cfg)
403
404
405
406
407
408
409
410


def get_ub(name: str):
    """Get userbuffer communicator corresponding to give key."""
    assert _ub_communicators is not None, "UB manager is not initialized."
    assert name in _ub_communicators, f"UB for {name} is not registered."
    return _ub_communicators[name]

411

412
413
414
415
416
417
418
def destroy_ub():
    """Destroy all allocated userbuffer communicators."""
    global _ub_communicators
    _ub_communicators = None
    global layers_atomic_ring_exchange
    layers_atomic_ring_exchange = []

419
420
421
422
423
424
425

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

    def __init__(self) -> None:
        super().__init__()
        assert torch.cuda.is_available(), "TransformerEngine needs CUDA."
426
        self.name = None
427
428
429
430
        self.fp8_initialized = False
        self.fp8 = False
        self.fp8_calibration = False
        self.fp8_meta = {}
431
        self.fp8_meta["fp8_checkpoint"] = False
432
433
        self.fp8_meta["fp8_group"] = None
        self.fp8_meta_tensors_initialized = False
434
        self.quantizers = {"scaling_fwd": {}, "scaling_bwd": {}}
435
436
437
        self.tp_group = None
        self.tp_size = 1
        self.sequence_parallel = False
438
439
        self.param_init_meta = {}
        self.primary_weights_in_fp8 = FP8GlobalStateManager.with_fp8_parameters()
440
        self.preserve_high_precision_init_val = FP8GlobalStateManager.with_high_precision_init_val()
441
442
        self.fsdp_wrapped = False
        self.fsdp_group = None
443
        self._fp8_workspaces: Dict[str, QuantizedTensor] = {}
444
445
        self.activation_dtype: Optional[torch.dtype] = None

446
447
448
        if not TEDebugState.debug_enabled:
            TEDebugState.initialize()

449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
    # Names of attributes that can be set quickly (see __setattr__
    # method)
    _fast_setattr_names: Set[str] = {
        "activation_dtype",
        "fp8",
        "fp8_initialized",
        "fp8_calibration",
        "fp8_parameters",
    }

    def __setattr__(self, name: str, value: Any) -> None:
        if name in TransformerEngineBaseModule._fast_setattr_names:
            # torch.nn.Module has a custom __setattr__ that handles
            # modules, parameters, and buffers. This is unnecessary
            # overhead when setting plain attrs.
            self.__dict__[name] = value
        else:
            # Default case
            super().__setattr__(name, value)
468

469
    def adjust_amax_history_length(self, length: int, fwd: Optional[bool] = None) -> None:
470
471
472
473
        """
        Delayed scaling only.

        Increase or decrease size of amax history based on given `length`.
474
475
476
477
478
479
480
481
482
483

        .. 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:
484
485
486
            if meta_key not in self.fp8_meta:
                # Handles non-parameter FP8 modules, e.g. DPA.
                continue
487
488
489
490
491
            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 = (
492
493
                    self.fp8_meta[meta_key].amax_history[:length].clone()
                )
494
495
496
497
498
499
            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)
                )

500
501
502
            # Update quantizers with new amax pointers.
            self.quantizers[meta_key] = self.fp8_meta[meta_key].make_quantizers()

503
504
            # Update the global buffers with new amax and history pointers.
            if FP8GlobalStateManager.get_buffer_info() in self.fp8_meta:
505
506
507
                fwd_pos, fwd_key, bwd_pos, bwd_key = self.fp8_meta[
                    FP8GlobalStateManager.get_buffer_info()
                ]
508
509
510
511
512
                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."
513
514
515
                        FP8GlobalStateManager.global_amax_buffer[buffer_key][pos] = self.fp8_meta[
                            meta_key
                        ].amax_history[0]
516
                        FP8GlobalStateManager.global_amax_history_buffer[buffer_key][pos] = (
517
518
                            self.fp8_meta[meta_key].amax_history
                        )
519

520
    def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
521
522
523
        """Init scales and amaxes for fwd | bwd."""
        fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"

524
        # Return early if recipe state matches recipe
525
        if self.fp8_meta_tensors_initialized:
526
527
528
529
530
531
            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
532
533
534
535
            if recipe.float8_current_scaling() and isinstance(
                recipe_state, Float8CurrentScalingRecipeState
            ):
                return
536
537
538
539
            if recipe.float8_block_scaling() and isinstance(
                recipe_state, Float8BlockScalingRecipeState
            ):
                return
540
541
542

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

545
546
547
548
549
        # Initialize recipe state and quantizers
        recipe_state = RecipeState.create(
            recipe,
            mode=("forward" if fwd else "backward"),
            num_quantizers=num_fp8_tensors,
550
551
        )

552
553
554
555
        self.fp8_meta[fp8_meta_tensor_key] = recipe_state
        self.quantizers[fp8_meta_tensor_key] = recipe_state.make_quantizers()

    def init_fp8_meta_tensors(self, recipe: Recipe) -> None:
556
        """Init scales and amaxes."""
557
558
559
        self.set_meta_tensor(True, recipe)
        self.set_meta_tensor(False, recipe)

560
561
        self.fp8_meta_tensors_initialized = True

562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
    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."""
577

578
579
580
581
582
        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_(
583
584
                        torch.zeros_like(self.fp8_meta[key].amax_history)
                    )
585
586
587
                else:
                    assert key in fp8_meta_tensors, "Cannot reset fp8 tensors."
                    self.fp8_meta[key].scale.copy_(fp8_meta_tensors[key][0])
588
                    self.fp8_meta[key].amax_history.copy_(fp8_meta_tensors[key][1])
589

590
591
592
593
        with torch.no_grad():
            reset("scaling_fwd")
            reset("scaling_bwd")

594
595
    def get_extra_state(self) -> torch.Tensor:
        """Save before checkpointing."""
596

597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
        # 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
627
        fp8_checkpoint = self.fp8_meta["fp8_checkpoint"] or self.fp8 or self.fp8_calibration
628
        if fp8_checkpoint:
629
630

            # Copy tensors to CPU and store
631
            state = {}
632
633
634
635
636
637
            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)
638
639

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

648
649
650
651
        # Serialize state into byte tensor
        torch.cuda.synchronize()
        state_serialized = bytearray(pickle.dumps(state))
        state_serialized = torch.frombuffer(state_serialized, dtype=torch.uint8)
652
        return state_serialized
653
654
655
656
657
658

    def set_extra_state(self, state: torch.Tensor) -> None:
        """Load previous state."""
        if state is None:
            return

659
        # Load state
660
        if isinstance(state, torch.Tensor):
661
            # Default format: byte tensor with pickled data
662
            state = pickle.loads(state.detach().cpu().numpy().tobytes())
663
        elif isinstance(state, io.BytesIO):
664
            # Deprecated format with io.BytesIO
665
            state.seek(0)
666
            state = torch.load(state, map_location="cuda")
667
668
        else:
            raise RuntimeError("Unsupported checkpoint format.")
669
670
671

        if state is None:
            return
672

673
        # Load extra items
674
        self.fp8_meta.update(state["extra_fp8_variables"])
675
        self.fp8_meta["recipe"] = state["recipe"]
676
677
678
        if "global_fp8_buffer_pos_fwd_recompute" in self.fp8_meta:
            del self.fp8_meta["global_fp8_buffer_pos_fwd_recompute"]

679
        # Initialize before loading
680
        self.init_fp8_meta_tensors(self.fp8_meta["recipe"])
681
682
683
684
685
686
687
688
689
690
691

        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
692
693
694
695
696
        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)
697
        torch.cuda.synchronize()
698
699
700
701
702
703
704
705
706

    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():
            self.activation_dtype = torch.get_autocast_gpu_dtype()
            return

        # All checks after this have already been performed once, thus skip
707
        if self.activation_dtype == inp.dtype:
708
709
            return

710
711
712
713
714
715
716
717
        dtype = inp.dtype
        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}"
                )
        self.activation_dtype = dtype
718
719

    def set_tensor_parallel_group(self, tp_group: Union[dist_group_type, None]) -> None:
720
721
722
723
724
725
726
727
728
        """
        Set the tensor parallel group for the given
        module before executing the forward pass.

        Parameters
        ----------
        tp_group : ProcessGroup, default = `None`
                  tensor parallel process group.
        """
729
730
731
        self.tp_group = tp_group
        self.tp_group_initialized = True

732
733
734
    def _get_fp8_params(self) -> Union[List[torch.Tensor], None]:
        """returns the FP8 weights."""
        fp8_params = []
735
        for param in self.parameters(recurse=False):
736
            if isinstance(param, QuantizedTensor) and param.requires_grad:
737
738
739
740
741
                fp8_params.append(param)
        if len(fp8_params) == 0:
            return None
        return fp8_params

742
743
    # This routine is shared across FP8 and FP8_calibration paths so should not actually
    # assume FP8 execution.
744
    def init_fp8_metadata(self, num_gemms: int = 1) -> None:
745
        """Initialize fp8 related metadata and tensors during fprop."""
746
        self.fp8_parameters = FP8GlobalStateManager.with_fp8_parameters()
747
748
        self.fp8 = FP8GlobalStateManager.is_fp8_enabled()
        self.fp8_calibration = FP8GlobalStateManager.is_fp8_calibration()
749
        fp8_enabled = self.fp8 or self.fp8_calibration
750
        self.fp8_meta["fp8_checkpoint"] = self.fp8 or self.fp8_calibration
751

752
        if self.fp8_parameters or fp8_enabled:
753
754
755
756
            if (
                self.fp8_initialized
                and FP8GlobalStateManager.get_fp8_recipe() == self.fp8_meta["recipe"]
            ):
757
                # FP8 init has already been run and recipe is the same, don't do anything.
758
                return
759
            self.fp8_meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe()
760
761
762
763
764
765
766
767
768
769
770
        else:
            # If fp8 isn't enabled, turn off and return.
            self.fp8_initialized = False
            return

        if self.fp8_parameters and not self.fp8_initialized:
            self.fp8_meta["num_gemms"] = num_gemms
            self.init_fp8_meta_tensors(self.fp8_meta["recipe"])

        if fp8_enabled:
            # Set FP8 and other FP8 metadata
771
            self.fp8_meta["num_gemms"] = num_gemms
772
            self.fp8_meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group()
773
774
775
776
777
778

            # Set FP8_MAX per tensor according to recipe
            self.fp8_meta["fp8_max_fwd"] = self.fp8_meta["recipe"].fp8_format.value.max_fwd
            self.fp8_meta["fp8_max_bwd"] = self.fp8_meta["recipe"].fp8_format.value.max_bwd

            # Allocate scales and amaxes
779
            self.init_fp8_meta_tensors(self.fp8_meta["recipe"])
780
            self.fp8_initialized = True
781
782

            self.fp8_meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe()
783
784
785
786
787
788

    @contextmanager
    def prepare_forward(
        self,
        inp: torch.Tensor,
        num_gemms: int = 1,
789
        allow_non_contiguous: bool = False,
Jan Bielak's avatar
Jan Bielak committed
790
    ) -> Generator[torch.Tensor, None, None]:
791
792
793
794
795
796
797
798
        """Checks and prep for FWD.
        The context manager is needed because there isn't a way for a module to know
        if it's the last FP8 module in the forward autocast. It is useful
        to setup the forward aggregated amax reduction for every module
        just in case. The autocast exit will pick up the most recent one.
        """
        # Activation recomputation is used and this is the second forward phase.
        if self.fp8 and in_fp8_activation_recompute_phase():
799
            FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta)
800
801
802
803
804
805
806
        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)
807
            self.init_fp8_metadata(num_gemms=num_gemms)
808

809
            if self.fp8 and self.sequence_parallel and self.fp8_meta["recipe"].delayed():
810
811
812
813
                assert self.fp8_meta["recipe"].reduce_amax, (
                    "Amax reduction across tensor parallel group is "
                    "necessary when using sequence parallelism with FP8."
                )
814

815
            if self.fp8 and not FP8GlobalStateManager.fp8_graph_capturing():
816
                FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(self.fp8_meta)
817
818

            # Activation recomputation is used and this is the first forward phase.
819
            if self.fp8 and self.training and is_fp8_activation_recompute_enabled():
820
                FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta)
821
822

        with torch.cuda.nvtx.range(self.__class__.__name__ + " forward"):
823
824
825
            if not allow_non_contiguous and not inp.is_contiguous():
                inp = inp.contiguous()
            yield inp
826
827

        if self.fp8 and in_fp8_activation_recompute_phase():
828
            FP8GlobalStateManager.restore_fp8_meta_tensors(self.fp8_meta)
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848

    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(
849
850
851
852
        ctx,
        grad_output: torch.Tensor,
        row_parallel_mode: bool,
        quantizer: Optional[Quantizer],
853
854
855
    ) -> Tuple[Union[torch.Tensor, None], ...]:
        """Utility function for backward.
        Returns tuple in order (all optional/None based on training precion/recipe):
856
857
            R1: gathered `grad_output`.
            R2: bias gradient on R1.
858
859

        """
860
861
        grad_output = grad_output.reshape((-1, grad_output.shape[-1]))
        grad_output = grad_output.contiguous()
862
863
        gather_grad_output = row_parallel_mode and ctx.sequence_parallel

864
        # Non-FP8 case: bgrad is fused with wgrad for this case.
865
        if not ctx.fp8 and not ctx.debug:
866
            if gather_grad_output:
867
                if not ctx.ub_overlap_ag:
868
                    grad_output, _ = gather_along_first_dim(grad_output, ctx.tp_group)
869
                else:
870
871
872
873
874
                    ctx.ub_obj_gradout.copy_into_buffer(grad_output, quantizer, local_chunk=True)
                    grad_output = ctx.ub_obj_gradout.get_buffer(quantizer)
            return grad_output, None

        # FP8 with all-gather: unfused bgrad, fused cast + transpose
875
        # Also supports debug quantization, which is handled inside gather_along_first_dim.
876
877
        if gather_grad_output:
            grad_bias = None
878
            if ctx.use_bias:
879
                grad_bias = grad_output.view(-1, grad_output.shape[-1]).sum(dim=0)
880
            if ctx.ub_overlap_ag:
881
882
                # Quantize the gradient if needed
                if not isinstance(
883
884
885
886
887
888
889
                    grad_output,
                    (
                        QuantizedTensor,
                        Float8TensorBase,
                        MXFP8TensorBase,
                        Float8BlockwiseQTensorBase,
                    ),
890
891
892
893
894
895
                ):
                    grad_output = quantizer(grad_output)

                # Copy into communication buffer, and replace original gradient with it
                ctx.ub_obj_gradout.copy_into_buffer(grad_output, quantizer, local_chunk=True)
                grad_output = ctx.ub_obj_gradout.get_buffer(quantizer)
896
            else:
897
898
899
900
                grad_output, _ = gather_along_first_dim(
                    grad_output,
                    ctx.tp_group,
                    quantizer=quantizer,
901
                )
902
            return grad_output, grad_bias
903

904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
        # 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)
            if (
                isinstance(
                    grad_output_.get_tensor(True),
                    (QuantizedTensor, Float8TensorBase, MXFP8TensorBase),
                )
                and ctx.use_bias
            ):
                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

921
922
        # FP8 without all-gather: fused bgrad + cast + transpose
        grad_bias = None
923
        if ctx.use_bias:
924
925
926
927
            if isinstance(
                grad_output,
                (QuantizedTensor, Float8TensorBase, MXFP8TensorBase, Float8BlockwiseQTensorBase),
            ):
928
                grad_bias = grad_output.dequantize().view(-1, grad_output.shape[-1]).sum(dim=0)
929
            else:
930
931
932
933
934
935
936
937
938
                if isinstance(quantizer, Float8BlockQuantizer):
                    # 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)
        if not isinstance(
            grad_output,
            (QuantizedTensor, Float8TensorBase, MXFP8TensorBase, Float8BlockwiseQTensorBase),
        ):
939
940
            grad_output = quantizer(grad_output)
        return grad_output, grad_bias
941

942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
    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)
        self.param_init_meta[name] = _ParameterInitMeta(**kwargs)

    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):
            # Ensure parameter is on a real device
961
962
            if param.device == torch.device("meta"):
                param = torch.empty_like(param, device="cuda")
963
964
965
966
967
968
969

            # 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:
970
971
972
973
974
975
                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)
976

977
            # If primary weights are in fp8, wrap the parameter as FP8Tensor
978
            fp8_meta_index = self.param_init_meta[name].fp8_meta_index
979
            high_precision_init_val = None
980
            if self.primary_weights_in_fp8 and fp8_meta_index is not None:
981
982
983
                if self.preserve_high_precision_init_val:
                    high_precision_init_val = param.detach().cpu()

984
985
986
987
988
989
                quantizer = self.quantizers["scaling_fwd"][fp8_meta_index]
                assert (
                    quantizer is not None
                )  # to use primary fp8 weight one needs to use FP8 autocast with specific recipe.
                quantizer.internal = False
                param = quantizer(param)
990
991
992
993
994

            # 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.
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
            param = torch.nn.Parameter(param)
            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)

            setattr(self, name, param)
1023

1024
1025
1026
    @abstractmethod
    def forward(self):
        """Needs override."""
1027

1028
    def get_weight_workspace(
1029
        self,
1030
1031
        *,
        tensor: Optional[torch.Tensor] = None,
1032
        quantizer: Optional[Quantizer] = None,
1033
1034
1035
        cache_name: Optional[str] = None,
        update_workspace: bool = True,
        skip_update_flag: Optional[torch.Tensor] = None,
1036
        fsdp_group: Optional[dist_group_type] = None,
1037
        workspace_dtype: Optional[torch.dtype] = None,
1038
    ) -> QuantizedTensor:
1039
1040
1041
1042
1043
1044
1045
1046
1047
        """Get FP8 workspace buffer and maybe update its values

        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.
1048
1049
1050
        quantizer: Quantizer, optional
            Quantizer used to cast the weights. Required if the
            workspace is being constructed or updated.
1051
1052
1053
1054
1055
1056
1057
        cache_name: str, optional
            Key for caching.
        update_workspace: bool, default = `True`
            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.
1058
1059
        fsdp_group: bool, default = None
            FSDP process group that the weights are distributed over.
1060
1061
1062
        workspace_dtype: torch.dtype, default = None
            If weight workspace contains high-precision tensor - for example
            for debug quantization, this is dtype of the tensor.
1063
1064
        """

1065
1066
1067
1068
1069
1070
1071
1072
1073
        # FP8 primary weights
        if isinstance(tensor, QuantizedTensor):
            if update_workspace and quantizer is not None:
                tensor.update_usage(
                    rowwise_usage=quantizer.rowwise_usage,
                    columnwise_usage=quantizer.columnwise_usage,
                )
            return tensor

1074
        # Try getting workspace from cache
1075
        out = None
1076

1077
1078
        if cache_name is not None:
            out = self._fp8_workspaces.get(cache_name, None)
1079
1080
1081
1082
1083
1084
1085
            if quantizer is not None and isinstance(out, MXFP8TensorBase):
                if quantizer.rowwise_usage and out._rowwise_data is None:
                    out = None
                    del self._fp8_workspaces[cache_name]
                elif quantizer.columnwise_usage and out._columnwise_data is None:
                    out = None
                    del self._fp8_workspaces[cache_name]
1086

1087
1088
1089
1090
1091
            is_debug = isinstance(quantizer, DebugQuantizer)
            is_out_debug_tensor = out is not None and isinstance(out, DebugQuantizedTensor)
            if is_debug != is_out_debug_tensor:
                out = None

1092
1093
1094
1095
1096
        # 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
1097
            and tensor is not None
1098
            and fsdp_group is not None
1099
            and out.data.shape != tensor.data.shape
1100
1101
1102
1103
        ):
            _fsdp_gather_tensors(fsdp_group, [tensor.data.shape], out)

        # Construct workspace if needed
1104
        if out is None:
1105
            if tensor is None or quantizer is None:
1106
                raise ValueError(
1107
                    "tensor and quantizer kwargs must be provided to construct FP8 workspace"
1108
                )
1109
1110
1111
1112
1113
1114
1115

            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
1116
            out = quantizer.quantize(tensor, dtype=workspace_dtype)
1117
1118
            if cache_name is not None:
                quantizer.internal = quantizer_internal
1119
1120

            # Update cache
1121
1122
            if cache_name is not None:
                self._fp8_workspaces[cache_name] = out
1123
            return out
1124
1125
1126
1127
1128
1129

        # Update workspace if needed
        if skip_update_flag is not None:
            update_workspace = True
        if update_workspace:
            if tensor is None:
1130
                raise ValueError("tensor kwarg must be provided to update FP8 workspace")
1131
            if hasattr(out, "quantize_"):
1132
                out.quantize_(tensor, noop_flag=skip_update_flag)
1133
1134
            else:
                tex.quantize(tensor, quantizer, out, skip_update_flag)
1135
        return out
1136

1137
1138
1139
    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
        """
        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])
1154
1155
1156
        super()._load_from_state_dict(
            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
        )
1157

1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
    def backward_dw(self):
        """
        Execute the delayed weight gradient computation.
        This method is called after the main backward pass to compute weight gradients.
        """
        if self.wgrad_store is None or not self.wgrad_store.delay_wgrad_compute():
            return
        with torch.cuda.nvtx.range(f"_{self.__class__.__name__}_wgrad"):
            (wgrad, grad_bias_, _, _), _ = self.wgrad_store.pop()
            if not self.fuse_wgrad_accumulation:
                unfused_weights = [getattr(self, name) for name in self.weight_names]
                weight_tensor = noop_cat(unfused_weights)
                if weight_tensor.grad is None:
                    weight_tensor.grad = wgrad.to(weight_tensor.dtype)
            if self.use_bias:
                bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names])
                if bias_tensor.grad is None:
                    bias_tensor.grad = grad_bias_.to(bias_tensor.dtype)
            del grad_bias_
            del wgrad

1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
    def _validate_name(self):
        """
        Validate name passed to the module.
        This is invoked in the forward() method as module names are assigned after Model is initialized in Megatron-LM.
        If no name is assigned, it creates a default name with layer count as the variable.
        """
        assert TEDebugState.debug_enabled
        import nvdlfw_inspect.api as debug_api

        if self.name is None:
            debug_api.log_message(
                "Names are not provided to debug modules. ",
                "Creating and using generic names. Pass names to debug modules for better"
                " insight. ",
                level=logging.WARNING,
            )
            self.name = f"Layer_{TEDebugState.get_layer_count()}"

    def _turn_off_unsupported_features_in_debug(self):
        if (
            getattr(self, "ub_bulk_wgrad", False)
            or getattr(self, "ub_bulk_dgrad", False)
            or getattr(self, "ub_overlap_ag", False)
            or getattr(self, "ub_overlap_rs_dgrad", False)
            or getattr(self, "ub_overlap_rs", False)
        ):
            import nvdlfw_inspect.api as debug_api

            debug_api.log_message(
                "UserBuffers are not supported in debug module. "
                "Using UB optimization will not affect the debug module. ",
                level=logging.WARNING,
            )
            if hasattr(self, "ub_bulk_wgrad"):
                self.ub_bulk_wgrad = None
            if hasattr(self, "ub_bulk_dgrad"):
                self.ub_bulk_dgrad = None
            if hasattr(self, "ub_overlap_ag"):
                self.ub_overlap_ag = None
            if hasattr(self, "ub_overlap_rs_dgrad"):
                self.ub_overlap_rs_dgrad = None
            if hasattr(self, "ub_overlap_rs"):
                self.ub_overlap_rs = None