base.py 52.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
yuguo's avatar
yuguo committed
46
from torch.utils.cpp_extension import IS_HIP_EXTENSION
47

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

50
51
52
_2X_ACC_FPROP = False
_2X_ACC_DGRAD = True
_2X_ACC_WGRAD = True
53
_multi_stream_cublas_workspace = []
54
_dummy_wgrads = {}
55
_multi_stream_cublas_batchgemm_workspace = []
56
57
_cublas_workspace = None
_ub_communicators = None
58
59
ub_stream_nums = int(os.getenv("NVTE_UB_STREAM_NUMS", "2"))
_NUM_MAX_UB_STREAMS = ub_stream_nums if IS_HIP_EXTENSION else 3
60
_MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY = None, None
61
layers_atomic_ring_exchange = []
62
63
64
65


def get_cublas_workspace_size_bytes() -> None:
    """Return 32 MiB if using hopper, 4 MiB for all other architectures."""
yuguo's avatar
yuguo committed
66
67
    # Add env for control the padding for blaslt
    if IS_HIP_EXTENSION:
68
        return 134_217_728
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
    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


84
85
86
87
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:
88
        for _ in range(tex._num_cublas_streams):
89
90
91
92
93
94
            _multi_stream_cublas_workspace.append(
                torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device="cuda")
            )
    return _multi_stream_cublas_workspace


yuguo's avatar
yuguo committed
95
96
97
98
99
100
def get_multi_stream_cublas_batchgemm_workspace() -> List[torch.Tensor]:
    """Returns workspace for multi-stream cublas."""
    global _multi_stream_cublas_batchgemm_workspace
    if not _multi_stream_cublas_batchgemm_workspace:
        for _ in range(tex._num_cublas_batchgemm_streams):
            _multi_stream_cublas_batchgemm_workspace.append(
yuguo's avatar
yuguo committed
101
                torch.empty(128, dtype=torch.uint8, device="cuda")
yuguo's avatar
yuguo committed
102
103
104
            )
    return _multi_stream_cublas_batchgemm_workspace

105

yuguo's avatar
yuguo committed
106
107
108
109
if bool(int(os.getenv("NVTE_DISABLE_FC2_DGRAD_OVERLAP", "0"))):
    remove_ag_gemm_dgrad = ["fc2_dgrad"]
else:
    remove_ag_gemm_dgrad = []
110

111

112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
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()


128
129
def initialize_ub(
    shape: list,
130
    tp_size: int,
131
    use_fp8: bool = False,
132
    dtype: torch.dtype = torch.bfloat16,
133
    ub_cfgs: Optional[dict] = None,
134
    bootstrap_backend: Union[str, torch.distributed.Backend] = None,
135
) -> None:
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
    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.
    """
183
    if not tex.device_supports_multicast():
yuguo's avatar
yuguo committed
184
        assert bool(int(os.getenv("UB_SKIPMC", "1"))), (
185
186
187
188
            "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."
        )

189
190
191
    global _ub_communicators
    assert _ub_communicators is None, "UB communicators are already initialized."
    _ub_communicators = {}
192
193

    if tex.ubuf_built_with_mpi():
194
195
        # 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...
196
        assert torch.distributed.is_mpi_available()
197
198
        _ = torch.distributed.new_group(backend="mpi")
        helper = tex.CommOverlapHelper()
199
    else:
200
201
        # Bootstrapping with torch.distributed API, so check backend and construct
        # intra/inter-node process groups...
202
203
204
205
206
        assert (
            torch.distributed.is_initialized()
        ), "torch.distributed must be initialized before Userbuffers"
        if bootstrap_backend is None:
            bootstrap_backend = "nccl"
207
            if torch.distributed.is_mpi_available():
208
                bootstrap_backend = "mpi"
209
210
            elif torch.distributed.is_gloo_available():
                bootstrap_backend = "gloo"
211
        else:
212
213
214
215
216
217
218
219
220
            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."
            )
221
222
223
224
225

        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)

226
227
        num_domains = world_size // tp_size
        mydomain_idx = world_rank // tp_size
228
        if num_domains > 1:
229
230
231
232
            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(
233
234
                ranks_per_domain_list, backend=bootstrap_backend
            )
235
236
            local_rank = torch.distributed.get_rank(tp_domain_group)
            tp_domain_ranks = torch.distributed.get_process_group_ranks(tp_domain_group)
237

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

            helper = tex.CommOverlapHelper(world_group)
246

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

256
    # Allocate cuBLAS workspace with expanded size for chunking in overlapping GEMM calls
257
    global _cublas_workspace
258
259
260
261
262
263
264
    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)
265
266

    # Default buffer precision: AllGather buffers use fp8 when using fp8 recipe
267
    layers_all_gather_overlap = [
268
269
270
271
272
273
        "qkv_fprop",
        "qkv_dgrad",
        "proj_dgrad",
        "fc1_fprop",
        "fc1_dgrad",
        "fc2_dgrad",
274
    ]
275
    layers_reduce_scatter_overlap = ["proj_fprop", "fc2_fprop", "qkv_wgrad", "fc1_wgrad"]
Jaemin Choi's avatar
Jaemin Choi committed
276
    dgrad_reduce_scatter_overlap = ["qkv_dgrad", "fc1_dgrad"]
277
    # Default overlap methods for layers
yuguo's avatar
yuguo committed
278
279
280
281
282
283
284
285
286
287
288
289
    if bool(int(os.getenv("NVTE_NO_PIPELINE_OVERLAP", "0"))):
        methods = {
            "ring_exchange": ["qkv_fprop", "fc1_fprop", "proj_dgrad", "fc2_dgrad", "proj_fprop", "fc2_fprop"],
            "pipeline": [],
            "bulk": ["qkv_dgrad", "qkv_wgrad", "fc1_dgrad", "fc1_wgrad"],
        }
    else:
        methods = {
            "ring_exchange": ["qkv_fprop", "fc1_fprop", "proj_dgrad", "fc2_dgrad"],
            "pipeline": ["proj_fprop", "fc2_fprop"],
            "bulk": ["qkv_dgrad", "qkv_wgrad", "fc1_dgrad", "fc1_wgrad"],
        }
290

291
    # AG-RS overlap pairs of layers forming a tensor-parallel block
292
293
    ag_rs_pairs = {"qkv_fprop": "proj_fprop", "fc1_fprop": "fc2_fprop"}
    rs_ag_pairs = {v: k for k, v in ag_rs_pairs.items()}
294
295
296
    global layers_atomic_ring_exchange
    layers_atomic_ring_exchange = []

297
298
299
300
301
302
    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.")

303
    def get_default_config(name):
304
        global _MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY
305
306
        method = get_method(name)
        is_reduce_scatter = name in layers_reduce_scatter_overlap
307
308
        if _MIN_STREAM_PRIORITY is None or _MAX_STREAM_PRIORITY is None:
            _MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY = tex.get_stream_priority_range()
309
310
311
        default_cfg = {
            "method": method,
            "is_reduce_scatter": is_reduce_scatter,
yuguo's avatar
yuguo committed
312
            "num_sm": 1 if method == "ring_exchange" else 8,
313
            "cga_size": 1 if method == "ring_exchange" else 2,
314
315
            "set_sm_margin": not method == "ring_exchange",
            "num_splits": tp_size if method == "ring_exchange" else 4,
316
317
318
319
            "aggregate": False,
            "atomic_gemm": False,
            "use_ce": True,
            "fp8_buf": name in layers_all_gather_overlap,
320
321
322
            "comm_priority": _MAX_STREAM_PRIORITY,
            "gemm_priority": _MIN_STREAM_PRIORITY,
            "pipeline_rs_overlap_first_gemm": False,
323
324
325
        }
        return default_cfg

326
327
328
    def add_ub(
        name: str,
        method: str,
329
        is_reduce_scatter: bool,
330
331
        num_sm: int = 16,
        cga_size: int = 2,
332
        set_sm_margin: bool = False,
333
        num_splits: int = 0,
334
335
        aggregate: bool = False,
        atomic_gemm: bool = False,
336
        use_ce: bool = True,
337
        fp8_buf: bool = False,
338
339
340
        comm_priority: int = 0,
        gemm_priority: int = 0,
        pipeline_rs_overlap_first_gemm: bool = False,
341
    ) -> None:
342
343
344
345
346
        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."
347
            if method == "bulk":
348
                warnings.warn(
349
                    f"At {name}, atoimic GEMM not is supported for a bulk overlap."
350
351
352
                    "Defaulting to `atomic_gemm=False`."
                )
                atomic_gemm = 0
353
        if not is_reduce_scatter and method == "pipeline":
354
            raise ValueError(
355
                f"At {name}, `pipeline` overlap method is not supported for AllGather."
356
            )
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
        # 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

376
        buffer_dtype = torch.uint8 if (use_fp8 and fp8_buf) else dtype
377
        if method == "ring_exchange":
378
379
380
381
            ub_obj = tex.CommOverlapP2P(
                shape,  # Communication buffer shape
                buffer_dtype,  # Communication buffer data type
                helper,  # Helper for torch.distributed callbacks during bootstrapping
382
                tp_size,  # Tensor-parallel group size (may be different than local_size)
383
384
385
386
387
388
389
390
                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,
391
392
                gemm_priority=gemm_priority,
                comm_priority=comm_priority,
393
            )
394
        else:
395
396
397
398
            ub_obj = tex.CommOverlap(
                shape,  # Communication buffer shape
                buffer_dtype,  # Communication buffer data type
                helper,  # Helper for torch.distributed callbacks during bootstrapping
399
                tp_size,  # Tensor-parallel group size (may be different than local_size)
400
                num_splits=num_splits,
yuguo's avatar
yuguo committed
401
                num_max_streams=_NUM_MAX_UB_STREAMS,
402
403
404
405
                comm_cga_size=cga_size,
                num_comm_sm=num_sm,
                set_sm_margin=set_sm_margin,
                atomic_gemm=atomic_gemm,
406
407
408
                gemm_priority=gemm_priority,
                comm_priority=comm_priority,
                rs_overlap_first_gemm=pipeline_rs_overlap_first_gemm,
409
            )
410
411
        _ub_communicators[name] = ub_obj

Jaemin Choi's avatar
Jaemin Choi committed
412
413
    if ub_cfgs is not None:
        for name in dgrad_reduce_scatter_overlap:
414
415
            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
416
417
                assert wgrad_name not in ub_cfgs
                layers_reduce_scatter_overlap.remove(wgrad_name)
418
                layers_all_gather_overlap.remove(name)
Jaemin Choi's avatar
Jaemin Choi committed
419
                layers_reduce_scatter_overlap.append(name)
420
421
422
                methods["bulk"].remove(name)
                new_method = ub_cfgs[name]["method"]
                methods[new_method].append(name)
Jaemin Choi's avatar
Jaemin Choi committed
423

424
    for name in methods["ring_exchange"] + methods["pipeline"] + methods["bulk"]:
yuguo's avatar
yuguo committed
425
426
        if name in remove_ag_gemm_dgrad:
            continue
427
        ub_cfg = get_default_config(name)
428
        if ub_cfgs is not None and name in ub_cfgs:
429
            fp8_buf = (name in layers_all_gather_overlap) or (
430
                ub_cfgs[name].get("fp8_buf", False) and name in methods["pipeline"]
431
            )
432
433
434
            ub_cfg.update(ub_cfgs[name])
            ub_cfg["fp8_buf"] = fp8_buf
        add_ub(name, **ub_cfg)
435
436
437
438
439


def get_ub(name: str):
    """Get userbuffer communicator corresponding to give key."""
    assert _ub_communicators is not None, "UB manager is not initialized."
yuguo's avatar
yuguo committed
440
441
442
    # assert name in _ub_communicators, f"UB for {name} is not registered."
    if name in remove_ag_gemm_dgrad:
        return None
443
444
    return _ub_communicators[name]

445

446
447
448
449
450
451
452
def destroy_ub():
    """Destroy all allocated userbuffer communicators."""
    global _ub_communicators
    _ub_communicators = None
    global layers_atomic_ring_exchange
    layers_atomic_ring_exchange = []

453
454
455
456
457
458
459

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

    def __init__(self) -> None:
        super().__init__()
        assert torch.cuda.is_available(), "TransformerEngine needs CUDA."
460
        self.name = None
461
462
463
464
        self.fp8_initialized = False
        self.fp8 = False
        self.fp8_calibration = False
        self.fp8_meta = {}
465
        self.fp8_meta["fp8_checkpoint"] = False
466
467
        self.fp8_meta["fp8_group"] = None
        self.fp8_meta_tensors_initialized = False
468
        self.quantizers = {"scaling_fwd": {}, "scaling_bwd": {}}
469
470
471
        self.tp_group = None
        self.tp_size = 1
        self.sequence_parallel = False
472
473
        self.param_init_meta = {}
        self.primary_weights_in_fp8 = FP8GlobalStateManager.with_fp8_parameters()
474
        self.preserve_high_precision_init_val = FP8GlobalStateManager.with_high_precision_init_val()
475
476
        self.fsdp_wrapped = False
        self.fsdp_group = None
477
        self._fp8_workspaces: Dict[str, QuantizedTensor] = {}
478
479
        self.activation_dtype: Optional[torch.dtype] = None

480
481
482
        if not TEDebugState.debug_enabled:
            TEDebugState.initialize()

483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
    # 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)
502

503
    def adjust_amax_history_length(self, length: int, fwd: Optional[bool] = None) -> None:
504
505
506
507
        """
        Delayed scaling only.

        Increase or decrease size of amax history based on given `length`.
508
509
510
511
512
513
514
515
516
517

        .. 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:
518
519
520
            if meta_key not in self.fp8_meta:
                # Handles non-parameter FP8 modules, e.g. DPA.
                continue
521
522
523
524
525
            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 = (
526
527
                    self.fp8_meta[meta_key].amax_history[:length].clone()
                )
528
529
530
531
532
533
            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)
                )

534
535
536
            # Update quantizers with new amax pointers.
            self.quantizers[meta_key] = self.fp8_meta[meta_key].make_quantizers()

537
538
            # Update the global buffers with new amax and history pointers.
            if FP8GlobalStateManager.get_buffer_info() in self.fp8_meta:
539
540
541
                fwd_pos, fwd_key, bwd_pos, bwd_key = self.fp8_meta[
                    FP8GlobalStateManager.get_buffer_info()
                ]
542
543
544
545
546
                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."
547
548
549
                        FP8GlobalStateManager.global_amax_buffer[buffer_key][pos] = self.fp8_meta[
                            meta_key
                        ].amax_history[0]
550
                        FP8GlobalStateManager.global_amax_history_buffer[buffer_key][pos] = (
551
552
                            self.fp8_meta[meta_key].amax_history
                        )
553

554
    def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
555
556
557
        """Init scales and amaxes for fwd | bwd."""
        fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"

558
        # Return early if recipe state matches recipe
559
        if self.fp8_meta_tensors_initialized:
560
561
562
563
564
565
            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
566
567
568
569
            if recipe.float8_current_scaling() and isinstance(
                recipe_state, Float8CurrentScalingRecipeState
            ):
                return
570
571
572
573
            if recipe.float8_block_scaling() and isinstance(
                recipe_state, Float8BlockScalingRecipeState
            ):
                return
574
575
576

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

579
580
581
582
583
        # Initialize recipe state and quantizers
        recipe_state = RecipeState.create(
            recipe,
            mode=("forward" if fwd else "backward"),
            num_quantizers=num_fp8_tensors,
584
585
        )

586
587
588
589
        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:
590
        """Init scales and amaxes."""
591
592
593
        self.set_meta_tensor(True, recipe)
        self.set_meta_tensor(False, recipe)

594
595
        self.fp8_meta_tensors_initialized = True

596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
    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."""
611

612
613
614
615
616
        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_(
617
618
                        torch.zeros_like(self.fp8_meta[key].amax_history)
                    )
619
620
621
                else:
                    assert key in fp8_meta_tensors, "Cannot reset fp8 tensors."
                    self.fp8_meta[key].scale.copy_(fp8_meta_tensors[key][0])
622
                    self.fp8_meta[key].amax_history.copy_(fp8_meta_tensors[key][1])
623

624
625
626
627
        with torch.no_grad():
            reset("scaling_fwd")
            reset("scaling_bwd")

628
629
    def get_extra_state(self) -> torch.Tensor:
        """Save before checkpointing."""
630

631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
        # 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
661
        fp8_checkpoint = self.fp8_meta["fp8_checkpoint"] or self.fp8 or self.fp8_calibration
662
        if fp8_checkpoint:
663
664

            # Copy tensors to CPU and store
665
            state = {}
666
667
668
669
670
671
            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)
672
673

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

682
683
684
685
        # Serialize state into byte tensor
        torch.cuda.synchronize()
        state_serialized = bytearray(pickle.dumps(state))
        state_serialized = torch.frombuffer(state_serialized, dtype=torch.uint8)
686
        return state_serialized
687
688
689
690
691
692

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

693
        # Load state
694
        if isinstance(state, torch.Tensor):
695
            # Default format: byte tensor with pickled data
696
            state = pickle.loads(state.detach().cpu().numpy().tobytes())
697
        elif isinstance(state, io.BytesIO):
698
            # Deprecated format with io.BytesIO
699
            state.seek(0)
700
            state = torch.load(state, map_location="cuda")
701
702
        else:
            raise RuntimeError("Unsupported checkpoint format.")
703
704
705

        if state is None:
            return
706

707
        # Load extra items
708
        self.fp8_meta.update(state["extra_fp8_variables"])
709
        self.fp8_meta["recipe"] = state["recipe"]
710
711
712
        if "global_fp8_buffer_pos_fwd_recompute" in self.fp8_meta:
            del self.fp8_meta["global_fp8_buffer_pos_fwd_recompute"]

713
        # Initialize before loading
714
        self.init_fp8_meta_tensors(self.fp8_meta["recipe"])
715
716
717
718
719
720
721
722
723
724
725

        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
726
727
728
729
730
        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)
731
        torch.cuda.synchronize()
732
733
734
735
736
737
738
739
740

    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
741
        if self.activation_dtype == inp.dtype:
742
743
            return

744
745
746
747
748
749
750
751
        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
752
753

    def set_tensor_parallel_group(self, tp_group: Union[dist_group_type, None]) -> None:
754
755
756
757
758
759
760
761
762
        """
        Set the tensor parallel group for the given
        module before executing the forward pass.

        Parameters
        ----------
        tp_group : ProcessGroup, default = `None`
                  tensor parallel process group.
        """
763
764
765
        self.tp_group = tp_group
        self.tp_group_initialized = True

766
767
768
    def _get_fp8_params(self) -> Union[List[torch.Tensor], None]:
        """returns the FP8 weights."""
        fp8_params = []
769
        for param in self.parameters(recurse=False):
770
            if isinstance(param, QuantizedTensor) and param.requires_grad:
771
772
773
774
775
                fp8_params.append(param)
        if len(fp8_params) == 0:
            return None
        return fp8_params

776
777
    # This routine is shared across FP8 and FP8_calibration paths so should not actually
    # assume FP8 execution.
778
    def init_fp8_metadata(self, num_gemms: int = 1) -> None:
779
        """Initialize fp8 related metadata and tensors during fprop."""
780
        self.fp8_parameters = FP8GlobalStateManager.with_fp8_parameters()
781
782
        self.fp8 = FP8GlobalStateManager.is_fp8_enabled()
        self.fp8_calibration = FP8GlobalStateManager.is_fp8_calibration()
783
        fp8_enabled = self.fp8 or self.fp8_calibration
784
        self.fp8_meta["fp8_checkpoint"] = self.fp8 or self.fp8_calibration
785

786
        if self.fp8_parameters or fp8_enabled:
787
788
789
790
            if (
                self.fp8_initialized
                and FP8GlobalStateManager.get_fp8_recipe() == self.fp8_meta["recipe"]
            ):
791
                # FP8 init has already been run and recipe is the same, don't do anything.
792
                return
793
            self.fp8_meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe()
794
795
796
797
798
799
800
801
802
803
804
        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
805
            self.fp8_meta["num_gemms"] = num_gemms
806
            self.fp8_meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group()
807
808
809
810
811
812

            # 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
813
            self.init_fp8_meta_tensors(self.fp8_meta["recipe"])
814
            self.fp8_initialized = True
815
816

            self.fp8_meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe()
817
818
819
820
821
822

    @contextmanager
    def prepare_forward(
        self,
        inp: torch.Tensor,
        num_gemms: int = 1,
823
        allow_non_contiguous: bool = False,
Jan Bielak's avatar
Jan Bielak committed
824
    ) -> Generator[torch.Tensor, None, None]:
825
826
827
828
829
830
831
832
        """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():
833
            FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta)
834
835
836
837
838
839
840
        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)
841
            self.init_fp8_metadata(num_gemms=num_gemms)
842

843
            if self.fp8 and self.sequence_parallel and self.fp8_meta["recipe"].delayed():
844
845
846
847
                assert self.fp8_meta["recipe"].reduce_amax, (
                    "Amax reduction across tensor parallel group is "
                    "necessary when using sequence parallelism with FP8."
                )
848

849
            if self.fp8 and not FP8GlobalStateManager.fp8_graph_capturing():
850
                FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(self.fp8_meta)
851
852

            # Activation recomputation is used and this is the first forward phase.
853
            if self.fp8 and self.training and is_fp8_activation_recompute_enabled():
854
                FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta)
855
856

        with torch.cuda.nvtx.range(self.__class__.__name__ + " forward"):
857
858
859
            if not allow_non_contiguous and not inp.is_contiguous():
                inp = inp.contiguous()
            yield inp
860
861

        if self.fp8 and in_fp8_activation_recompute_phase():
862
            FP8GlobalStateManager.restore_fp8_meta_tensors(self.fp8_meta)
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882

    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(
883
884
885
886
        ctx,
        grad_output: torch.Tensor,
        row_parallel_mode: bool,
        quantizer: Optional[Quantizer],
887
888
889
    ) -> Tuple[Union[torch.Tensor, None], ...]:
        """Utility function for backward.
        Returns tuple in order (all optional/None based on training precion/recipe):
890
891
            R1: gathered `grad_output`.
            R2: bias gradient on R1.
892
893

        """
894
895
        grad_output = grad_output.reshape((-1, grad_output.shape[-1]))
        grad_output = grad_output.contiguous()
896
897
        gather_grad_output = row_parallel_mode and ctx.sequence_parallel

898
        # Non-FP8 case: bgrad is fused with wgrad for this case.
899
        if not ctx.fp8 and not ctx.debug:
900
            if gather_grad_output:
yuguo's avatar
yuguo committed
901
                if not ctx.ub_overlap_ag or ctx.ub_obj_gradout is None:
902
                    grad_output, _ = gather_along_first_dim(grad_output, ctx.tp_group)
903
                else:
904
905
906
907
908
                    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
909
        # Also supports debug quantization, which is handled inside gather_along_first_dim.
910
911
        if gather_grad_output:
            grad_bias = None
912
            if ctx.use_bias:
913
                grad_bias = grad_output.view(-1, grad_output.shape[-1]).sum(dim=0)
yuguo's avatar
yuguo committed
914
            if ctx.ub_overlap_ag and ctx.ub_obj_gradout is not None:
915
916
                # Quantize the gradient if needed
                if not isinstance(
917
918
919
920
921
922
923
                    grad_output,
                    (
                        QuantizedTensor,
                        Float8TensorBase,
                        MXFP8TensorBase,
                        Float8BlockwiseQTensorBase,
                    ),
924
925
926
927
928
929
                ):
                    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)
930
            else:
931
932
933
934
                grad_output, _ = gather_along_first_dim(
                    grad_output,
                    ctx.tp_group,
                    quantizer=quantizer,
935
                )
936
            return grad_output, grad_bias
937

938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
        # 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

955
956
        # FP8 without all-gather: fused bgrad + cast + transpose
        grad_bias = None
957
        if ctx.use_bias:
958
959
960
961
            if isinstance(
                grad_output,
                (QuantizedTensor, Float8TensorBase, MXFP8TensorBase, Float8BlockwiseQTensorBase),
            ):
962
                grad_bias = grad_output.dequantize().view(-1, grad_output.shape[-1]).sum(dim=0)
963
            else:
964
965
966
967
968
969
970
971
972
                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),
        ):
973
974
            grad_output = quantizer(grad_output)
        return grad_output, grad_bias
975

976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
    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
995
996
            if param.device == torch.device("meta"):
                param = torch.empty_like(param, device="cuda")
997
998
999
1000
1001
1002
1003

            # 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:
1004
1005
1006
1007
1008
1009
                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)
1010

1011
            # If primary weights are in fp8, wrap the parameter as FP8Tensor
1012
            fp8_meta_index = self.param_init_meta[name].fp8_meta_index
1013
            high_precision_init_val = None
1014
            if self.primary_weights_in_fp8 and fp8_meta_index is not None:
1015
1016
1017
                if self.preserve_high_precision_init_val:
                    high_precision_init_val = param.detach().cpu()

1018
1019
1020
1021
1022
1023
                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)
1024
1025
1026
1027
1028

            # 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.
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
            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)
1057

1058
1059
1060
    @abstractmethod
    def forward(self):
        """Needs override."""
1061

1062
    def get_weight_workspace(
1063
        self,
1064
1065
        *,
        tensor: Optional[torch.Tensor] = None,
1066
        quantizer: Optional[Quantizer] = None,
1067
1068
1069
        cache_name: Optional[str] = None,
        update_workspace: bool = True,
        skip_update_flag: Optional[torch.Tensor] = None,
1070
        fsdp_group: Optional[dist_group_type] = None,
1071
        workspace_dtype: Optional[torch.dtype] = None,
1072
    ) -> QuantizedTensor:
1073
1074
1075
1076
1077
1078
1079
1080
1081
        """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.
1082
1083
1084
        quantizer: Quantizer, optional
            Quantizer used to cast the weights. Required if the
            workspace is being constructed or updated.
1085
1086
1087
1088
1089
1090
1091
        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.
1092
1093
        fsdp_group: bool, default = None
            FSDP process group that the weights are distributed over.
1094
1095
1096
        workspace_dtype: torch.dtype, default = None
            If weight workspace contains high-precision tensor - for example
            for debug quantization, this is dtype of the tensor.
1097
1098
        """

1099
1100
1101
1102
1103
1104
1105
1106
1107
        # 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

1108
        # Try getting workspace from cache
1109
        out = None
1110

1111
1112
        if cache_name is not None:
            out = self._fp8_workspaces.get(cache_name, None)
1113
1114
1115
1116
1117
1118
1119
            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]
1120

1121
1122
1123
1124
1125
            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

1126
1127
1128
1129
1130
        # 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
1131
            and tensor is not None
1132
            and fsdp_group is not None
1133
            and out.data.shape != tensor.data.shape
1134
1135
1136
1137
        ):
            _fsdp_gather_tensors(fsdp_group, [tensor.data.shape], out)

        # Construct workspace if needed
1138
        if out is None:
1139
            if tensor is None or quantizer is None:
1140
                raise ValueError(
1141
                    "tensor and quantizer kwargs must be provided to construct FP8 workspace"
1142
                )
1143
            out = quantizer.quantize(tensor, dtype=workspace_dtype)
1144
1145

            # Update cache
1146
1147
            if cache_name is not None:
                self._fp8_workspaces[cache_name] = out
1148
            return out
1149
1150
1151
1152
1153
1154

        # Update workspace if needed
        if skip_update_flag is not None:
            update_workspace = True
        if update_workspace:
            if tensor is None:
1155
                raise ValueError("tensor kwarg must be provided to update FP8 workspace")
1156
            if hasattr(out, "quantize_"):
1157
                out.quantize_(tensor, noop_flag=skip_update_flag)
1158
1159
            else:
                tex.quantize(tensor, quantizer, out, skip_update_flag)
1160
        return out
1161

1162
1163
1164
    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
        """
        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])
1179
1180
1181
        super()._load_from_state_dict(
            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
        )
1182

1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
    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

1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
    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