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

"""Linear API"""
6
from typing import Callable, Dict, Optional, Tuple, Union, List
7
8
from functools import reduce
from operator import mul as multiply_op
9
import warnings
10
import os
11
12
13

import torch

14
import transformer_engine_torch as tex
15

16
from transformer_engine.common.recipe import Recipe
Paweł Gadziński's avatar
Paweł Gadziński committed
17
from transformer_engine.pytorch.torch_version import torch_version
18

19
from .base import (
20
21
    fill_userbuffers_buffer_for_all_gather,
    get_dummy_wgrad,
22
23
24
25
26
27
    get_ub,
    TransformerEngineBaseModule,
    _2X_ACC_FPROP,
    _2X_ACC_DGRAD,
    _2X_ACC_WGRAD,
)
28
from ._common import noop_cat, WeightGradStore
29
from ..quantization import FP8GlobalStateManager
30
31
from ..utils import (
    cast_if_needed,
32
    clear_tensor_data,
33
    divide,
34
    init_method_constant,
35
36
    requires_grad,
    needs_quantized_gemm,
37
    assert_dim_for_fp8_exec,
38
    assert_dim_for_all_gather,
39
40
    nvtx_range_pop,
    nvtx_range_push,
41
    get_nvtx_range_context,
42
43
44
45
46
)
from ..distributed import (
    set_tensor_model_parallel_attributes,
    get_distributed_world_size,
    allreduce,
47
    symmetric_all_reduce,
48
49
    reduce_scatter_along_first_dim,
    gather_along_first_dim,
50
    is_fp8_activation_recompute_enabled,
51
    in_fp8_activation_recompute_phase,
52
53
    _fsdp_scatter_tensors,
    _fsdp_gather_tensors,
54
55
)
from ..cpp_extensions import (
56
    general_gemm,
57
)
58
from ..constants import GemmParallelModes, dist_group_type
59
from ..jit import no_torch_dynamo
60
from ..graph import is_graph_capturing
61
from ..quantized_tensor import (
62
    QuantizedTensor,
63
    QuantizedTensorStorage,
64
65
66
67
    Quantizer,
    prepare_for_saving,
    restore_from_saved,
)
68
from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer
69
from ..tensor.mxfp8_tensor import MXFP8Quantizer
70
from ..tensor.utils import is_custom
71
from ..export import is_in_onnx_export_mode, assert_warmed_up
72
73
74
75
76
77
from ..cpu_offload import (
    is_cpu_offload_enabled,
    start_offload,
    mark_not_offload,
    mark_activation_offload,
)
78
from ...debug.pytorch.debug_state import TEDebugState
79

80
81
82
83
84
85
86
87
88
89
90
__all__ = ["Linear"]


class _Linear(torch.autograd.Function):
    """Linear semi-top level module
    Calls custom cuda extensions.
    """

    @staticmethod
    def forward(
        ctx,
91
        weight: torch.Tensor,
92
        inp: torch.Tensor,
93
        bias: Optional[torch.Tensor],
94
        non_tensor_args: Tuple,
95
    ) -> torch.Tensor:
96
        # pylint: disable=missing-function-docstring
97

98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
        (
            is_first_microbatch,
            fp8,
            fp8_calibration,
            wgrad_store,
            input_quantizer,
            weight_quantizer,
            output_quantizer,
            grad_input_quantizer,
            grad_weight_quantizer,
            grad_output_quantizer,
            fuse_wgrad_accumulation,
            cpu_offloading,
            tp_group,
            tp_size,
            sequence_parallel,
            tensor_parallel,
            activation_dtype,
            parallel_mode,
            is_grad_enabled,
            ub_overlap_rs_fprop,
            ub_overlap_ag_dgrad,
            ub_overlap_ag_fprop,
            ub_overlap_rs_dgrad,
            ub_bulk_dgrad,
            ub_bulk_wgrad,
            ub_name,
            fp8_output,  # pylint: disable=unused-variable
            fsdp_group,
            module,
            skip_fp8_weight_update,
            symmetric_ar_type,
            save_original_input,
            debug,
        ) = non_tensor_args

134
135
136
137
138
        # NVTX label for profiling
        nvtx_label = "transformer_engine._Linear.forward"
        if ub_name is not None:
            nvtx_label = f"{nvtx_label}.{ub_name}"

139
        # Make sure input dimensions are compatible
140
        out_features, in_features = weight.shape
141
        assert inp.shape[-1] == in_features, "GEMM not possible"
142

143
        # Configure tensor-parallel communication
144
        tp_world_size = get_distributed_world_size(tp_group)
145
146
147
148
        backward_needs_input = is_grad_enabled and weight.requires_grad
        with_input_all_gather_nccl = (
            parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop
        )
149
150

        # Configure Userbuffers communication (comm+GEMM overlap)
151
152
153
154
155
156
        if debug:  # turn off userbuffers in debug mode
            ub_overlap_rs_fprop = False
            ub_overlap_ag_fprop = False
            ub_overlap_rs_dgrad = False
            ub_bulk_wgrad = False
            ub_bulk_dgrad = False
157
158
159
        ub_obj = None
        ub_type = None
        if ub_overlap_rs_fprop:
160
            ub_obj = get_ub(ub_name + "_fprop", fp8)
161
162
            ub_type = tex.CommOverlapType.RS
        elif ub_overlap_ag_fprop:
163
            ub_obj = get_ub(ub_name + "_fprop", fp8)
164
165
            ub_type = tex.CommOverlapType.AG

166
167
        # custom recipe check
        custom = is_custom(input_quantizer) or is_custom(weight_quantizer)
168

169
170
171
172
173
174
175
176
        # ------------------------------------------------------
        # Prepare input tensor
        # Note: Cast to expected dtype and perform tensor-parallel communication
        # ------------------------------------------------------
        nvtx_range_push(f"{nvtx_label}.input_cast_comm")
        inputmat = inp  # Input tensor to save for backward (maybe sharded)
        inputmat_total = None  # Input tensor to pass to GEMM (gathered)
        own_quantized_input = False
177
        if fp8:
178
            assert_dim_for_fp8_exec(inputmat, weight)
179
            assert_dim_for_all_gather(inputmat, with_input_all_gather_nccl, input_quantizer)
180
181
182
183
184
            if save_original_input:
                assert not isinstance(
                    input_quantizer, Float8Quantizer
                ), "DelayedScaling recipe is not supported with save_original_input"

185
186
187
188
189
190
        if with_input_all_gather_nccl or ub_overlap_ag_fprop:  # All-gather input tensor

            # Cast local input tensor if needed
            if fp8 or debug:
                if input_quantizer is None:
                    raise ValueError("Missing quantizer for input tensor")
191
                if not isinstance(inputmat, QuantizedTensorStorage) and not custom:
192
193
                    own_quantized_input = True
                    input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input)
194
195
196
197
198
                    if isinstance(
                        input_quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer)
                    ):
                        # All-gather is not supported with FP8 column-wise data
                        input_quantizer.set_usage(columnwise=False)
199
200
201
202
203
                    if save_original_input:
                        # No need for column-wise data since this
                        # tensor will not be cached for backward pass
                        input_quantizer.set_usage(columnwise=False)
                        own_quantized_input = False
204
                    inputmat = input_quantizer(inputmat)
205
            else:
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
                inputmat = cast_if_needed(inp, activation_dtype)  # Cast for AMP

            # Initialize gathered input tensor
            quantizer = None
            if fp8 or debug:
                quantizer = input_quantizer
                quantizer.set_usage(rowwise=True, columnwise=False)
            if with_input_all_gather_nccl:  # Perform NCCL all-gather
                inputmat_total, _ = gather_along_first_dim(
                    inputmat,
                    tp_group,
                    quantizer=quantizer,
                )
            elif ub_overlap_ag_fprop:  # Initialize Userbuffers all-gather
                inputmat_total, _ = fill_userbuffers_buffer_for_all_gather(
                    ub_obj,
                    inputmat,
                    quantizer,
                    tp_group,
                )

        else:  # Do not all-gather input tensor
            if fp8 or debug:
229
                if isinstance(inputmat, QuantizedTensorStorage):
230
                    inputmat.update_usage(rowwise_usage=True)
231
                else:
232
233
                    if input_quantizer is None:
                        raise ValueError("Missing quantizer for input tensor")
234
235
236
                    input_quantizer.set_usage(
                        rowwise=True, columnwise=backward_needs_input and not save_original_input
                    )
237
                    inputmat = input_quantizer(inputmat)
238
                    own_quantized_input = True
239
            else:
240
241
                inputmat = cast_if_needed(inp, activation_dtype)  # Cast for AMP
            inputmat_total = inputmat
242
243
244

        if is_cpu_offload_enabled():
            start_offload(inputmat)
245
        nvtx_range_pop(f"{nvtx_label}.input_cast_comm")
246
247
248
        # ------------------------------------------------------
        # Input tensor is ready for GEMM...
        # ------------------------------------------------------
249

250
251
252
        # ------------------------------------------------------
        # Prepare weight tensor
        # ------------------------------------------------------
253
254
        weightmat = weight
        if fp8 or debug:
255
            # Configure quantizer
256
257
            # No need to set the quantizer states if weight is already quantized
            if weight_quantizer is not None and not isinstance(weight, QuantizedTensor):
258
259
260
261
262
263
264
                columnwise_usage = is_grad_enabled and inp.requires_grad
                if not columnwise_usage:
                    columnwise_usage = (
                        is_fp8_activation_recompute_enabled()
                        and not in_fp8_activation_recompute_phase()
                    )
                weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)
265
266
267
            elif isinstance(weight, QuantizedTensor):
                # If weight is already quantized, no need to set quantizer states
                weight_quantizer = weight._quantizer
268
            # Get quantized weight
269
270
271
272
273
274
275
276
            update_workspace = is_first_microbatch is None or is_first_microbatch
            weightmat = module.get_weight_workspace(
                tensor=weight,
                quantizer=weight_quantizer,
                cache_name=(None if is_first_microbatch is None else "weight"),
                update_workspace=update_workspace,
                skip_update_flag=skip_fp8_weight_update,
                fsdp_group=fsdp_group,
277
                workspace_dtype=activation_dtype,
278
            )
279
280
            weightmat.update_usage(rowwise_usage=True)

281
        else:
282
283
284
285
            weightmat = cast_if_needed(weightmat, activation_dtype)  # Cast for AMP
        # ------------------------------------------------------
        # Weight tensor is ready for GEMM...
        # ------------------------------------------------------
286
287
288

        # Cast bias to expected dtype
        bias_dtype = activation_dtype
289
        if needs_quantized_gemm(inputmat_total) and activation_dtype == torch.float32:
290
            # cuBLAS does not support FP8 GEMM with FP32 bias, so we cast to BF16
291
292
293
294
295
296
297
298
299
300
            bias_dtype = torch.bfloat16
        bias = cast_if_needed(bias, bias_dtype) if bias is not None else bias

        # Calibrate quantizers if needed
        if not fp8 and fp8_calibration:
            if input_quantizer is not None:
                input_quantizer.calibrate(inputmat_total)
            if weight_quantizer is not None:
                weight_quantizer.calibrate(weight)

301
302
        # Choose whether to use GEMM kernel with split accumulator
        use_split_accumulator = _2X_ACC_FPROP
303
304
305
        if fp8:
            recipe = FP8GlobalStateManager.get_fp8_recipe()
            if hasattr(recipe, "fp8_gemm_fprop"):
306
307
308
309
310
                use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator

        # Configure output quantizer
        if output_quantizer is not None:
            output_quantizer.set_usage(rowwise=True, columnwise=False)
311

312
313
314
315
316
317
318
319
320
321
322
323
324
325
        # Output buffer for Userbuffers reduce-scatter
        reduce_scatter_out = None
        if ub_overlap_rs_fprop:
            out_shape = list(inp.shape)
            out_shape[0] //= tp_world_size
            out_shape[-1] = out_features
            reduce_scatter_out = torch.empty(out_shape, dtype=activation_dtype, device=inp.device)

        # ------------------------------------------------------
        # Forward GEMM
        # Note: y = x * w^T
        # ------------------------------------------------------
        nvtx_range_push(f"{nvtx_label}.gemm")
        gemm_out, *_, reduce_scatter_out = general_gemm(
326
327
328
            weightmat,
            inputmat_total,
            quantization_params=output_quantizer,
329
            out_dtype=activation_dtype,
330
            bias=bias,
331
            use_split_accumulator=use_split_accumulator,
332
333
            ub=ub_obj,
            ub_type=ub_type,
334
            extra_output=reduce_scatter_out,
335
        )
336
        nvtx_range_pop(f"{nvtx_label}.gemm")
337
338
339
340
        # ------------------------------------------------------
        # Finished forward GEMM...
        # ------------------------------------------------------

341
342
343
344
345
346
347
        # Deallocate GEMM input tensor if no longer needed
        # TODO(yuzhongw, tmoon): Figure out why inputmat_total is not automatically
        # deallocated by GC. Manually deallocating is a temporary hack.
        if with_input_all_gather_nccl:
            clear_tensor_data(inputmat_total)
            inputmat_total = None

348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
        # ------------------------------------------------------
        # Prepare output tensor
        # Note: Perform tensor-parallel communication
        # ------------------------------------------------------
        out = None
        if ub_overlap_rs_fprop:
            out = reduce_scatter_out
        elif parallel_mode == "row" and tp_size > 1:
            nvtx_range_push(f"{nvtx_label}.row_parallel_comm")
            out = gemm_out
            if sequence_parallel:
                out, _ = reduce_scatter_along_first_dim(out, tp_group)
            elif tensor_parallel:
                if symmetric_ar_type is not None:
                    out, _ = symmetric_all_reduce(out, tp_group, all_reduce_type=symmetric_ar_type)
                else:
                    out, _ = allreduce(out, tp_group)
            nvtx_range_pop(f"{nvtx_label}.row_parallel_comm")
        else:
            out = gemm_out
        # ------------------------------------------------------
        # Output tensor is ready to return...
        # ------------------------------------------------------

        # ------------------------------------------------------
        # Cache state for backward pass
        # ------------------------------------------------------
375
376

        if is_grad_enabled:
377
378
379
            if save_original_input:
                inputmat = inp

380
            ctx.weight_quantizer = weight_quantizer
381
382
383
384
385

            ctx.backward_input_needs_gather = (
                weight.requires_grad and parallel_mode == "column" and sequence_parallel
            )

386
387
388
389
            # Discard unneeded data in input tensor
            if (
                backward_needs_input
                and own_quantized_input
390
                and isinstance(inputmat, QuantizedTensorStorage)
391
            ):
392
393
394
                if (
                    ctx.backward_input_needs_gather
                    and weight_quantizer.supports_only_rowwise_all_gather()
395
396
397
398
399
400
401
402
403
                ):
                    # All-gather is not supported with FP8 column-wise data
                    inputmat.update_usage(rowwise_usage=True, columnwise_usage=False)
                else:
                    # Discard row-wise data since it is not needed in backward pass
                    inputmat.update_usage(rowwise_usage=False, columnwise_usage=True)

            # Cached input tensor
            saved_inputmat = None
404
405
            if backward_needs_input:
                saved_inputmat = inputmat
406

407
408
            if cpu_offloading and saved_inputmat is not None:
                mark_activation_offload(saved_inputmat)
409

410
411
            # Scatter intermediate/activation tensors saved for the backward pass
            # NOTE: FSDP sharding is not valid for models initialized with primary Fp8 weights
412
            nvtx_range_push(f"{nvtx_label}.fsdp_scatter")
413
414
415
            ctx.fsdp_group = fsdp_group
            ctx.fsdp_shapes = _fsdp_scatter_tensors(
                fsdp_group,
416
                saved_inputmat,
417
                weightmat if fp8 and not isinstance(weight, QuantizedTensorStorage) else None,
418
            )
419
            nvtx_range_pop(f"{nvtx_label}.fsdp_scatter")
420

evt_fugx1's avatar
evt_fugx1 committed
421
            if cpu_offloading or int(os.getenv("NVTE_SWAP_OVERLAP_GRAD", "0")):
dongchl's avatar
dongchl committed
422
                ctx.grad_added_to_main_grad = hasattr(weight, "grad_added_to_main_grad")
423

dongchl's avatar
dongchl committed
424
                if ctx.grad_added_to_main_grad:
425
426
427
428
429
430
431
                    # If you are passing torch.nn.Parameter through the Torch hooks, you will
                    # get back torch.Tensor. Torch rips off the Parameter wrapper.
                    # You need to preserve the weight object to have all the attributes user
                    # sets for the weights. Because of this, it is not recommended to offload
                    # weights if weights are externally touched outside this module
                    ctx.weight_object = weight

432
433
            if cpu_offloading:
                mark_not_offload(weight, weightmat, bias)
434
435
            # TODO(ksivamani): Check memory usage
            tensors_to_save, tensor_objects = prepare_for_saving(
436
                saved_inputmat,
437
                weightmat,
438
                weight,
439
                bias,
440
            )
441
442
            ctx.save_for_backward(*tensors_to_save)
            ctx.tensor_objects = tensor_objects
443

444
445
            ctx.activation_dtype = activation_dtype
            ctx.fp8 = fp8
446
            ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None
447
448
            ctx.input_quantizer = input_quantizer
            ctx.grad_input_quantizer = grad_input_quantizer
449
450
            ctx.grad_weight_quantizer = grad_weight_quantizer
            ctx.grad_output_quantizer = grad_output_quantizer
451
            ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation
452
            if fuse_wgrad_accumulation and weight.requires_grad:
453
454
455
456
457
458
459
460
                # This check is needed to ensure that main_grad is not created
                # during the forward pass when using MCore FSDP as it creates
                # the main_grad buffer lazily before backprop
                if hasattr(weight, "__fsdp_param__"):
                    # MCore FSDP creates main_grad lazily before backward
                    ctx.main_grad_func = weight.get_main_grad
                else:
                    ctx.main_grad_func = lambda: weight.main_grad
461

462
            ctx.debug = debug
463
            ctx.custom = custom
464
            ctx.cpu_offloading = cpu_offloading
465
            ctx.is_first_microbatch = is_first_microbatch
466
            ctx.use_bias = bias is not None
467
468
            ctx.sequence_parallel = sequence_parallel
            ctx.tensor_parallel = tensor_parallel
469
            ctx.inp_shape = inp.shape
470
471
            ctx.parallel_mode = parallel_mode
            ctx.tp_group = tp_group
472
473
474
475
            ctx.ub_overlap_ag = ub_overlap_ag_dgrad
            ctx.ub_overlap_rs_dgrad = ub_overlap_rs_dgrad
            ctx.ub_bulk_dgrad = ub_bulk_dgrad
            ctx.ub_bulk_wgrad = ub_bulk_wgrad
476
            ctx.ub_name = ub_name
477
478
            ctx.tp_size = tp_size
            ctx.requires_dgrad = inp.requires_grad
479
            ctx.requires_wgrad = weight.requires_grad
480
            ctx.reduce_and_update_bwd_fp8_tensors = False
481

482
            ctx.owns_input = saved_inputmat is not inp
483
            if ctx.fp8 and requires_grad(inp, weight, bias):
484
485
486
487
                _first_fp8_module = FP8GlobalStateManager.IS_FIRST_FP8_MODULE
                ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module()
                if in_fp8_activation_recompute_phase():
                    FP8GlobalStateManager.IS_FIRST_FP8_MODULE = _first_fp8_module
488
            ctx.wgrad_store = wgrad_store
489

490
491
492
        # ------------------------------------------------------
        # Cached state for backward pass is ready...
        # ------------------------------------------------------
493

494
        return out
495
496

    @staticmethod
497
    def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
498
        # pylint: disable=missing-function-docstring
499

500
501
502
503
504
        # NVTX label for profiling
        nvtx_label = "transformer_engine._Linear.backward"
        if ctx.ub_name is not None:
            nvtx_label = f"{nvtx_label}.{ctx.ub_name}"

505
        with get_nvtx_range_context("_Linear_backward"):
506
507
508
509
            saved_tensors = ctx.saved_tensors
            inputmat, weight_fp8, weight, bias = (  # pylint: disable=unbalanced-tuple-unpacking
                restore_from_saved(ctx.tensor_objects, saved_tensors)
            )
510

511
512
513
            # Delete the references to tensor objects once they've been consumed
            # by the `restore_from_saved` method to construct back the actual tensors.
            ctx.tensor_objects = None
514
515
516

            # Since main_grad can be modified inplace, it should not be a part of saved_tensors
            main_grad = (
517
                ctx.main_grad_func()
518
519
520
521
                if weight is not None and ctx.fuse_wgrad_accumulation and ctx.requires_wgrad
                else None
            )

dongchl's avatar
dongchl committed
522
523
            if ctx.cpu_offloading or int(os.getenv("NVTE_SWAP_OVERLAP_GRAD", "0")):
                if ctx.grad_added_to_main_grad:
524
                    weight = ctx.weight_object
dongchl's avatar
dongchl committed
525

526
527
            if ctx.requires_wgrad and ctx.fuse_wgrad_accumulation:
                weight.main_grad = main_grad
528

529
530
531
            # Gather intermediate/activation tensors if needed
            # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already
            #       shards/unshards the base weights so we don't do it ourselves
532
            nvtx_range_push(f"{nvtx_label}.fsdp_gather")
533
534
535
536
            _fsdp_gather_tensors(
                ctx.fsdp_group,
                ctx.fsdp_shapes,
                inputmat,
537
                weight_fp8,
538
            )
539
            nvtx_range_pop(f"{nvtx_label}.fsdp_gather")
540

541
            # Configure Userbuffers communication (comm+GEMM overlap)
542
            ctx.ub_obj_gradout = None
543
            ub_obj_dgrad = None
544
            ub_obj_wgrad = None
545
546
            ub_type_dgrad = None
            ub_type_wgrad = None
547
            dgrad_shape = [reduce(multiply_op, ctx.inp_shape[:-1]), ctx.inp_shape[-1]]
548
            if ctx.ub_overlap_ag:
549
                # Overlap grad_output all-gather with dgrad compute
550
                ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8)
551
552
                ub_obj_dgrad = ctx.ub_obj_gradout
                ub_type_dgrad = tex.CommOverlapType.AG
553
554
            elif ctx.ub_overlap_rs_dgrad:
                # Overlap dgrad reduce-scatter with dgrad compute
555
                ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8)
556
557
                ub_obj_dgrad = ctx.ub_obj_gradout
                ub_type_dgrad = tex.CommOverlapType.RS
558
559
560
            else:
                if ctx.ub_bulk_dgrad:
                    # Overlap inputmat all-gather with dgrad compute
561
                    ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8)
562
563
                    ub_obj_dgrad = ctx.ub_obj_gradout
                    ub_type_dgrad = tex.CommOverlapType.AG
564
565
                if ctx.ub_bulk_wgrad:
                    # Overlap dgrad reduce-scatter with wgrad compute
566
                    ub_obj_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8)
567
                    ub_type_wgrad = tex.CommOverlapType.RS
568
569
570
571
572
573
574
575

            # --------------------------------------------------
            # Prepare grad output tensor
            # Note: Cast to expected dtype and perform tensor-parallel communication
            # --------------------------------------------------

            # Unmodified grad output tensor
            grad_output_arg = grad_output
576

577
578
579
580
            # Configure quantizer for grad output tensor
            # Note: dgrad GEMM requires row-wise usage, wgrad GEMM
            # requires column-wise usage
            if ctx.grad_output_quantizer is not None:
581
582
583
584
585
586
587
                quantizer = ctx.grad_output_quantizer
                quantizer.set_usage(rowwise=True, columnwise=True)
                if ctx.ub_overlap_ag:
                    # Userbuffers only supports communication for one
                    # tensor usage at a time. Configure quantizer with
                    # usage for only dgrad GEMM.
                    quantizer.set_usage(columnwise=False)
588

589
590
591
592
593
594
595
596
597
598
599
600
601
            # Adjust the quantization direction approach depending
            # on whether wgrad calculations will be performed.
            # NOTE: If requires_dgrad is False, disabling `rowwise` quantization and keeping `columnwise` quantization
            #       results in `Assertion failed: output_tensor->has_data(). Quantizing in only the columnwise direction not supported yet!`
            # NOTE: For `ctx.bias is True`, selected quantize kernel errors with
            #       `cast_kernels.cuh:1322 in function fp8_quantize_arch_l_100: Not implemented scaling mode or fusion: NVTE_DELAYED_TENSOR_SCALING or IS_DBIAS=true on GPU with compute capability < 10.0.`
            if (
                not ctx.use_bias
                and not ctx.requires_wgrad
                and ctx.grad_output_quantizer is not None
            ):
                ctx.grad_output_quantizer.set_usage(columnwise=False)

602
            # Prepare grad output tensor
603
            nvtx_range_push(f"{nvtx_label}.grad_output_preprocess")
604
605
606
607
            (
                grad_output,
                grad_bias,
            ) = TransformerEngineBaseModule.grad_output_preprocess(
608
609
610
611
                ctx,
                grad_output,
                ctx.parallel_mode == "row",
                ctx.grad_output_quantizer,
612
            )
613
            nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess")
614

615
616
617
618
619
620
621
622
623
624
            # --------------------------------------------------
            # Grad output tensor is ready for computing grad input...
            # --------------------------------------------------

            # --------------------------------------------------
            # Prepare input tensor
            # Note: Input tensor is needed for wgrad GEMM.
            # Tensor-parallel communication is overlapped with dgrad
            # GEMM.
            # --------------------------------------------------
625
            inputmat_total = None
626
            inputmat_total_work = None
627
628
            if ctx.requires_wgrad:
                if ctx.fp8 or ctx.debug:
629
                    if isinstance(inputmat, QuantizedTensorStorage):
630
631
                        # Input tensor is already quantized
                        pass
632
                    elif ctx.debug or ctx.custom:
633
634
635
636
                        # Debug quantizer will be applied immediately before wgrad GEMM
                        pass
                    else:
                        # Quantize input tensor
637
                        quantizer = ctx.input_quantizer
638
                        if quantizer.supports_only_rowwise_all_gather():
639
                            # All-gather is not supported with FP8 column-wise data
640
641
642
643
644
645
646
647
                            quantizer.set_usage(
                                rowwise=True,
                                columnwise=not ctx.backward_input_needs_gather,
                            )
                        else:
                            quantizer.set_usage(rowwise=False, columnwise=True)
                        inputmat = quantizer(inputmat)
                else:
648
                    if isinstance(inputmat, QuantizedTensorStorage):
649
650
651
                        inputmat = inputmat.dequantize(dtype=ctx.activation_dtype)
                    else:
                        inputmat = cast_if_needed(inputmat, ctx.activation_dtype)
652
            if ctx.backward_input_needs_gather:
653
                quantizer = None
654
                if ctx.fp8 or ctx.debug:
655
                    quantizer = ctx.input_quantizer
656
                    if quantizer.supports_only_rowwise_all_gather():
657
658
659
660
661
                        # If data is in FP8, we compute FP8 transposes manually
                        quantizer.set_usage(rowwise=True, columnwise=False)
                    else:
                        # wgrad GEMM requires input with column-wise usage
                        quantizer.set_usage(rowwise=False, columnwise=True)
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
                if ctx.ub_bulk_dgrad:
                    inputmat_total, _ = fill_userbuffers_buffer_for_all_gather(
                        ub_obj_dgrad,
                        inputmat,
                        quantizer,
                        ctx.tp_group,
                    )
                else:
                    nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input")
                    inputmat_total, inputmat_total_work = gather_along_first_dim(
                        inputmat,
                        ctx.tp_group,
                        async_op=True,
                        quantizer=quantizer,
                    )
                    nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input")
678
679
            else:
                inputmat_total = inputmat
680
681
682
            # --------------------------------------------------
            # Input tensor is ready for computing grad weight...
            # --------------------------------------------------
683

684
            # --------------------------------------------------
685
            # Compute grad input tensor
686
687
            # --------------------------------------------------

688
689
            dgrad = None
            dgrad_work = None
690
            if ctx.requires_dgrad:
691

692
                # Make sure required data is available
693
                if isinstance(grad_output, QuantizedTensorStorage):
694
                    grad_output.update_usage(rowwise_usage=True)
695
696
697
                if ctx.weight_quantizer is not None and isinstance(
                    weight_fp8, QuantizedTensorStorage
                ):
698
699
700
701
                    weight_fp8.update_usage(columnwise_usage=True)

                # Choose whether to use GEMM kernel with split accumulator
                use_split_accumulator = _2X_ACC_DGRAD
702
703
704
                if ctx.fp8:
                    recipe = ctx.fp8_recipe
                    if hasattr(recipe, "fp8_gemm_dgrad"):
705
                        use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator
706

707
708
709
710
711
712
713
714
715
716
                # Update grad input quantizer
                if ctx.grad_input_quantizer is not None:
                    ctx.grad_input_quantizer.set_usage(rowwise=True, columnwise=False)

                # Output buffers for Userbuffers reduce-scatter
                gemm_out = None
                reduce_scatter_out = None
                if ctx.ub_overlap_rs_dgrad:
                    reduce_scatter_out = torch.empty(
                        dgrad_shape, dtype=ctx.activation_dtype, device=grad_output_arg.device
717
                    )
718
719
                elif ctx.ub_bulk_wgrad:
                    gemm_out = ub_obj_wgrad.get_buffer(local_chunk=False)
720

721
722
                # dgrad GEMM
                # Note: dx = dy * w
723

724
725
                nvtx_range_push(f"{nvtx_label}.dgrad_gemm")
                gemm_out, *_, reduce_scatter_out = general_gemm(
726
727
728
729
730
                    weight_fp8,
                    grad_output,
                    layout="NN",
                    grad=True,
                    quantization_params=ctx.grad_input_quantizer,
731
                    out=gemm_out,
732
                    out_dtype=ctx.activation_dtype,
733
                    use_split_accumulator=use_split_accumulator,
734
735
                    ub=ub_obj_dgrad,
                    ub_type=ub_type_dgrad,
736
                    extra_output=reduce_scatter_out,
737
738
                    bulk_overlap=ctx.ub_bulk_dgrad,
                )
739
                nvtx_range_pop(f"{nvtx_label}.dgrad_gemm")
740

741
742
                # Prepare grad input tensor
                # Note: Perform tensor-parallel communication
743
                if ctx.ub_overlap_rs_dgrad:
744
745
746
747
                    dgrad = reduce_scatter_out
                elif ctx.ub_bulk_wgrad:
                    dgrad = ub_obj_wgrad.get_buffer(local_chunk=True)
                elif ctx.parallel_mode == "column" and ctx.tp_size > 1:
748
                    nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad")
749
                    dgrad = gemm_out
750
751
752
753
754
                    if ctx.sequence_parallel:
                        dgrad, dgrad_work = reduce_scatter_along_first_dim(
                            dgrad,
                            ctx.tp_group,
                            async_op=True,
755
                        )
756
                    else:
757
                        dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True)
758
                    nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad")
759
760
761
762
763
764
765
766
767
768
                else:
                    dgrad = gemm_out

            # --------------------------------------------------
            # Grad input tensor has been computed...
            # --------------------------------------------------

            # --------------------------------------------------
            # Compute grad weight
            # --------------------------------------------------
769

770
771
            wgrad = None
            if ctx.requires_wgrad:
772

773
774
775
                # Prepare input tensor
                # Note: Synchronize tensor-parallel communication and
                # make sure required data is available
776
777
778
                if inputmat_total_work is not None:
                    inputmat_total_work.wait()
                    inputmat_total_work = None
779
                if ctx.fp8 or ctx.debug:
780
                    if isinstance(inputmat_total, QuantizedTensorStorage):
781
782
783
784
785
786
787
788
789
                        inputmat_total.update_usage(columnwise_usage=True)
                    else:
                        ctx.input_quantizer.set_usage(rowwise=False, columnwise=True)
                        inputmat_total = ctx.input_quantizer(inputmat_total)

                # Prepare grad output tensor
                # Note: Synchronize tensor-parallel communication and
                # make sure required data is available
                if ctx.ub_overlap_ag and isinstance(ctx.grad_output_quantizer, MXFP8Quantizer):
790
                    # UB does not support pipelined overlapping grad output
791
792
793
                    # all-gather with wgrad GEMM. Also, we can't
                    # convert row-scaled MXFP8 to column-scaled, so we
                    # can't reuse the grad output that was gathered
794
                    # for the dgrad GEMM. We work around by explicitly
795
796
797
798
799
800
                    # overlapping the AG operation with the dgrad GEMM.

                    # Get the communication stream from the dgrad GEMM to use for the AG
                    dgrad_send_stream, dgrad_recv_stream = ub_obj_dgrad.get_communication_stream()

                    # This object is separate from the ub_obj_wgrad object which is passed to the GEMM
801
                    ub_obj_overlap_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8)
802

803
                    ctx.grad_output_quantizer.set_usage(rowwise=False, columnwise=True)
804
805
806
807
808
809
810

                    # We use the send stream to copy into the userbuffers.
                    # This is the same stream that we will use to access the data in the AG,
                    # so we dont need to add any syncs yet.
                    with torch.cuda.stream(dgrad_send_stream):
                        grad_output, _ = fill_userbuffers_buffer_for_all_gather(
                            ub_obj_overlap_wgrad,
811
                            grad_output_arg,
812
                            ctx.grad_output_quantizer,
813
814
                            ctx.tp_group,
                        )
815
816
817
818
819

                    # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm
                    tex.bulk_overlap_ag_with_external_gemm(
                        ub_obj_overlap_wgrad, dgrad_send_stream, dgrad_recv_stream
                    )
820

821
                if ctx.fp8 or ctx.debug:
822
                    if isinstance(grad_output, QuantizedTensorStorage):
823
824
825
826
                        grad_output.update_usage(columnwise_usage=True)
                    else:
                        ctx.grad_output_quantizer.set_usage(rowwise=False, columnwise=True)
                        grad_output = ctx.grad_output_quantizer(grad_output)
827

828
829
830
831
832
833
834
                # Figure out whether to use split accumulator
                use_split_accumulator = _2X_ACC_WGRAD
                if ctx.fp8:
                    recipe = ctx.fp8_recipe
                    if hasattr(recipe, "fp8_gemm_wgrad"):
                        use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator

835
836
837
838
839
840
841
842
843
                # Figure out whether to output wgrad GEMM directly into main grad
                if ctx.is_first_microbatch is not None:
                    accumulate_wgrad_into_param_main_grad = (
                        ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch
                    )
                else:
                    accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation

                # Output buffer for overlapping FP8 grad input
844
                # reduce-scatter with wgrad GEMM
845
                reduce_scatter_out = None
846
                if ctx.ub_bulk_wgrad and ub_obj_wgrad.is_fp8_ubuf():
847
848
                    reduce_scatter_out = torch.empty(
                        dgrad_shape, dtype=ctx.activation_dtype, device=grad_output_arg.device
849
850
                    )

851
852
853
                # Arguments to include in wgrad GEMM closure
                wgrad_gemm_kwargs = {
                    "out_dtype": (
854
855
                        main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype
                    ),
856
                    "quantization_params": ctx.grad_weight_quantizer,
857
858
859
860
861
                    "accumulate": (
                        accumulate_wgrad_into_param_main_grad
                        if not getattr(weight, "overwrite_main_grad", False)
                        else False
                    ),
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
                    "layout": "NT",
                    "out": main_grad if ctx.fuse_wgrad_accumulation else None,
                    "bias": (bias if (grad_bias is None and not ctx.fp8) else None),
                    "use_split_accumulator": use_split_accumulator,
                    "grad": True,
                    "ub": ub_obj_wgrad,
                    "ub_type": ub_type_wgrad,
                    "extra_output": reduce_scatter_out,
                    "bulk_overlap": ctx.ub_bulk_wgrad,
                }

                def wgrad_gemm(
                    x: torch.Tensor,
                    dy: torch.Tensor,
                ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
                    """Perform wgrad GEMM: dw = dy^T * x

                    May be fused with bgrad computation.

                    May be called outside of this function to enable
                    some advanced communication/compute overlapping.

                    """
                    nvtx_range_push(f"{nvtx_label}.wgrad_gemm")
                    dw, db, *_ = general_gemm(x, dy, **wgrad_gemm_kwargs)
                    nvtx_range_pop(f"{nvtx_label}.wgrad_gemm")
                    return dw, db

                # Choose whether to call wgrad GEMM now or delay
891
                if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute():
892
893
894
895
896
897
898
899
900
901
902
                    if (
                        wgrad_gemm_kwargs["ub"] is not None
                        or wgrad_gemm_kwargs["ub_type"] is not None
                        or wgrad_gemm_kwargs["extra_output"] is not None
                        or wgrad_gemm_kwargs["bulk_overlap"]
                    ):
                        raise NotImplementedError(
                            "Delayed weight grad computation is not supported "
                            "with Userbuffers (tensor-parallel communication overlapping)"
                        )
                    ctx.wgrad_store.put([inputmat_total, grad_output], wgrad_gemm)
903
904
905
                    # overlap_grad_reduce, dongcl
                    if int(os.getenv("NVTE_OVERLAP_GRAD_REDUCE", "0")):
                        wgrad = torch.empty(weight.size(), dtype=ctx.activation_dtype, device=weight.device)
906
907
                else:

908
909
910
911
                    # Call wgrad GEMM now
                    wgrad, grad_bias_ = wgrad_gemm(inputmat_total, grad_output)

                    # Update grad bias if needed
912
913
914
915
                    if grad_bias is None:
                        grad_bias = grad_bias_
                    del grad_bias_

916
                    # Deallocate tensors if permitted
917
                    if ctx.owns_input:
918
                        # Input tensor is internal
919
                        clear_tensor_data(inputmat_total)
920
921
                    elif ctx.backward_input_needs_gather:
                        # Gathered input tensor is internal
922
                        clear_tensor_data(inputmat_total)
923
924
925
                    if ctx.parallel_mode == "row" and ctx.sequence_parallel:
                        # Gathered grad output tensor is internal
                        clear_tensor_data(grad_output)
926

927
                # Update grad input if overlapping reduce-scatter with wgrad GEMM
928
929
                if ctx.ub_bulk_wgrad:
                    if ub_obj_wgrad.is_fp8_ubuf():
930
                        dgrad = reduce_scatter_out
931
                    else:
932
933
934
935
936
                        dgrad = ub_obj_wgrad.get_buffer(local_chunk=True).clone()

            # --------------------------------------------------
            # Grad weight has been computed...
            # --------------------------------------------------
937

938
            # Don't return grad bias if not needed
939
940
941
            if not ctx.use_bias:
                grad_bias = None

942
            # Make sure all tensor-parallel communication is finished
943
944
945
946
947
948
949
950
            if inputmat_total_work is not None:
                inputmat_total_work.wait()
                inputmat_total_work = None
            if dgrad_work is not None:
                dgrad_work.wait()
                dgrad_work = None

        if ctx.requires_wgrad:
951
            # Handle custom DDP from mcore.
952
953
954
955
956
            if (
                ctx.fuse_wgrad_accumulation
                and weight is not None
                and hasattr(weight, "grad_added_to_main_grad")
            ):
957
                weight.grad_added_to_main_grad = True
958
                if getattr(weight, "zero_out_wgrad", False):
959
960
961
962
                    wgrad = get_dummy_wgrad(
                        list(weight.main_grad.shape),
                        weight.dtype,
                        zero=True,
963
                    )
964
                else:
965
966
967
                    wgrad = get_dummy_wgrad(
                        list(weight.main_grad.shape),
                        weight.dtype,
968
                    )
969
970
971
972
            elif ctx.fuse_wgrad_accumulation:
                wgrad = None
        else:
            wgrad = None
973

974
        # Update FP8 scaling factors if needed
975
        if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing():
976
            nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors")
977
            FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
978
            nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors")
979

980
        # Scatter fp8 weight buffers
981
        if ctx.fp8 and not isinstance(weight, QuantizedTensorStorage):
982
            _fsdp_scatter_tensors(ctx.fsdp_group, weight_fp8)
983
        return (
984
            wgrad,
985
986
            dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None,
            grad_bias,
987
            None,
988
989
990
991
        )


class Linear(TransformerEngineBaseModule):
992
    """Applies a linear transformation to the incoming data :math:`y = xA^T + b`
993

Paweł Gadziński's avatar
Paweł Gadziński committed
994
    On NVIDIA GPUs it is a drop-in replacement for ``torch.nn.Linear``.
995
996
997
998
999
1000
1001

    Parameters
    ----------
    in_features : int
                 size of each input sample.
    out_features : int
                  size of each output sample.
Paweł Gadziński's avatar
Paweł Gadziński committed
1002
1003
1004
1005
1006
1007
    bias : bool, default = True
          if set to ``False``, the layer will not learn an additive bias.
    init_method : Callable, default = None
                 used for initializing weights in the following way: ``init_method(weight)``.
                 When set to ``None``, defaults to ``torch.nn.init.normal_(mean=0.0, std=0.023)``.
    get_rng_state_tracker : Callable, default = None
1008
                 used to get the random number generator state tracker for initializing weights.
Paweł Gadziński's avatar
Paweł Gadziński committed
1009
    rng_tracker_name : str, default = None
1010
                 the param passed to get_rng_state_tracker to get the specific rng tracker.
cyanguwa's avatar
cyanguwa committed
1011
    parameters_split : Optional[Union[Tuple[str, ...], Dict[str, int]]], default = None
1012
1013
1014
1015
1016
                      Configuration for splitting the weight and bias tensors along dim 0 into
                      multiple PyTorch parameters. If a list or tuple of strings is provided,
                      they are used to make the names of equally-sized parameters. If a dict
                      (preferably an OrderedDict) is provided, the keys are used as names and
                      values as split sizes along dim 0. The resulting parameters will have
Paweł Gadziński's avatar
Paweł Gadziński committed
1017
                      names that end in ``_weight`` or ``_bias``, so trailing underscores are
1018
                      stripped from any provided names.
1019
    device : Union[torch.device, str], default = "cuda"
1020
          The device on which the parameters of the model will be allocated. It is the user's
1021
1022
          responsibility to ensure all parameters are moved to the GPU before running the
          forward pass.
Paweł Gadziński's avatar
Paweł Gadziński committed
1023
    name : str, default = None
1024
        name of the module, currently used for debugging purposes.
1025
1026
1027

    Parallelism parameters
    ----------------------
Paweł Gadziński's avatar
Paweł Gadziński committed
1028
1029
1030
    sequence_parallel : bool, default = False
                       if set to ``True``, uses sequence parallelism.
    tp_group : ProcessGroup, default = None
1031
1032
1033
1034
              tensor parallel process group.
    tp_size : int, default = 1
             used as TP (tensor parallel) world size when TP groups are not formed during
             initialization. In this case, users must call the
Paweł Gadziński's avatar
Paweł Gadziński committed
1035
             ``set_tensor_parallel_group(tp_group)`` method on the initialized module before the
1036
1037
             forward pass to supply the tensor parallel group needed for tensor and sequence
             parallel collectives.
Paweł Gadziński's avatar
Paweł Gadziński committed
1038
    parallel_mode : {None, 'column', 'row'}, default = None
1039
1040
                   used to decide whether this Linear layer is Column Parallel Linear or Row
                   Parallel Linear as described `here <https://arxiv.org/pdf/1909.08053.pdf>`_.
Paweł Gadziński's avatar
Paweł Gadziński committed
1041
                   When set to ``None``, no communication is performed.
1042
1043
1044
1045

    Optimization parameters
    -----------------------
    fuse_wgrad_accumulation : bool, default = 'False'
Paweł Gadziński's avatar
Paweł Gadziński committed
1046
                             if set to ``True``, enables fusing of creation and accumulation of
1047
                             the weight gradient. When enabled, it is assumed that the weights
Paweł Gadziński's avatar
Paweł Gadziński committed
1048
1049
                             have an additional ``main_grad`` attribute (used instead of the
                             regular ``grad``) which is a pre-allocated buffer of the correct
1050
1051
                             size to accumulate gradients in. This argument along with
                             weight tensor having attribute 'overwrite_main_grad' set to True
Paweł Gadziński's avatar
Paweł Gadziński committed
1052
1053
1054
                             will overwrite ``main_grad`` instead of accumulating.
    return_bias : bool, default = False
                 when set to ``True``, this module will not apply the additive bias itself, but
1055
1056
1057
                 instead return the bias value during the forward pass together with the
                 output of the linear transformation :math:`y = xA^T`. This is useful when
                 the bias addition can be fused to subsequent operations.
Paweł Gadziński's avatar
Paweł Gadziński committed
1058
    params_dtype : torch.dtype, default = torch.get_default_dtype()
1059
1060
1061
                  it controls the type used to allocate the initial parameters. Useful when
                  the model is trained with lower precision and the original FP32 parameters
                  would not fit in GPU memory.
Paweł Gadziński's avatar
Paweł Gadziński committed
1062
1063
1064
    delay_wgrad_compute : bool, default = False
                         Whether or not to delay weight gradient computation. If set to ``True``,
                         it's the user's responsibility to call ``module.backward_dw`` to compute
1065
                         weight gradients.
1066
1067
1068
    symmetric_ar_type : {None, 'multimem_all_reduce', 'two_shot', 'one_shot'}, default = None
                   Type of symmetric memory all-reduce to use during the forward pass.
                   This can help in latency bound communication situations.
Paweł Gadziński's avatar
Paweł Gadziński committed
1069
                   Requires PyTorch version 2.7.0 or higher. When set to ``None``, standard all-reduce
1070
                   is used.
Paweł Gadziński's avatar
Paweł Gadziński committed
1071
1072
    save_original_input : bool, default = False
                       If set to ``True``, always saves the original input tensor rather than the
1073
1074
1075
                       cast tensor. In some scenarios, the input tensor is used by multiple modules,
                       and saving the original input tensor may reduce the memory usage.
                       Cannot work with FP8 DelayedScaling recipe.
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        sequence_parallel: bool = False,
        fuse_wgrad_accumulation: bool = False,
        tp_group: Optional[dist_group_type] = None,
        tp_size: int = 1,
        get_rng_state_tracker: Optional[Callable] = None,
1087
        rng_tracker_name: Optional[str] = None,
1088
1089
1090
        init_method: Optional[Callable] = None,
        bias: bool = True,
        return_bias: bool = False,
1091
        params_dtype: Optional[torch.dtype] = None,
1092
        parallel_mode: Optional[str] = None,
cyanguwa's avatar
cyanguwa committed
1093
        parameters_split: Optional[Union[Tuple[str, ...], Dict[str, int]]] = None,
1094
        device: Union[torch.device, str] = "cuda",
1095
        ub_overlap_ag: bool = False,
1096
        ub_overlap_rs: bool = False,
1097
        ub_overlap_rs_dgrad: bool = False,
1098
1099
        ub_bulk_dgrad: bool = False,
        ub_bulk_wgrad: bool = False,
1100
        ub_name: Optional[str] = None,
1101
        delay_wgrad_compute: bool = False,
1102
        symmetric_ar_type: Optional[str] = None,
1103
        save_original_input: bool = False,
1104
        name: Optional[str] = None,
1105
1106
    ) -> None:
        super().__init__()
1107
1108

        params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype
1109
1110
1111
1112
1113
1114
        self.in_features = in_features
        self.out_features = out_features
        self.fuse_wgrad_accumulation = fuse_wgrad_accumulation
        self.use_bias = bias
        self.return_bias = return_bias
        self.apply_bias = bias and not return_bias
1115
1116
        self.get_rng_state_tracker = get_rng_state_tracker
        self.rng_tracker_name = rng_tracker_name
1117
        self.symmetric_ar_type = symmetric_ar_type
1118
        self.save_original_input = save_original_input
1119
1120
        self.name = name

1121
1122
        self.wgrad_store = WeightGradStore(delay_wgrad_compute, ub_bulk_wgrad)

1123
1124
        if device == "meta":
            assert parameters_split is None, "Cannot split module parameters on 'meta' device."
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
        if tp_group is None:
            self.tp_size = tp_size
            if tp_size == 1:
                self.set_tensor_parallel_group(tp_group)
        else:
            self.tp_size = get_distributed_world_size(tp_group)
            self.set_tensor_parallel_group(tp_group)
        self.set_nccl_overlap_warning_if_tp()

        self.parallel_mode = parallel_mode
        assert (
            self.parallel_mode in GemmParallelModes
        ), f"parallel_mode {parallel_mode} not supported"

        if self.parallel_mode == "column":
            self.out_features = divide(self.out_features, self.tp_size)
        elif self.parallel_mode == "row":
            self.in_features = divide(self.in_features, self.tp_size)

        self.sequence_parallel = (self.tp_size > 1) and sequence_parallel

1146
        # Column parallel TP overlap options
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
        self.ub_overlap_ag_fprop = (
            self.parallel_mode == "column" and self.sequence_parallel and ub_overlap_ag
        )
        self.ub_overlap_rs_dgrad = (
            self.parallel_mode == "column" and self.sequence_parallel and ub_overlap_rs_dgrad
        )
        self.ub_bulk_dgrad = (
            self.parallel_mode == "column"
            and self.sequence_parallel
            and ub_bulk_dgrad
            and not self.ub_overlap_rs_dgrad
        )
        self.ub_bulk_wgrad = (
            self.parallel_mode == "column"
            and self.sequence_parallel
            and ub_bulk_wgrad
            and not self.ub_overlap_rs_dgrad
        )
1165
1166

        # Row parallel TP overlap options
1167
1168
1169
1170
1171
1172
        self.ub_overlap_rs_fprop = (
            self.parallel_mode == "row" and self.sequence_parallel and ub_overlap_rs
        )
        self.ub_overlap_ag_dgrad = (
            self.parallel_mode == "row" and self.sequence_parallel and ub_overlap_ag
        )
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186

        if any(
            [
                self.ub_overlap_rs_fprop,
                self.ub_overlap_ag_dgrad,
                self.ub_overlap_ag_fprop,
                self.ub_overlap_rs_dgrad,
                self.ub_bulk_dgrad,
                self.ub_bulk_wgrad,
            ]
        ):
            assert ub_name is not None, f"Comm+GEMM overlap layer '{ub_name}' is not initialized."
        self.ub_name = ub_name

1187
1188
1189
1190
1191
1192
1193
        if self.symmetric_ar_type is not None:
            assert torch_version() >= (
                2,
                7,
                0,
            ), "Torch version must be at least 2.7 to use symmetric memory"

1194
1195
1196
        # Initialize params in FP8
        with_fp8_params = FP8GlobalStateManager.with_fp8_parameters()

1197
1198
1199
1200
1201
1202
1203
1204
        # Contiguous buffers for params
        weight_tensor = torch.empty(
            self.out_features,
            self.in_features,
            device=device,
            dtype=params_dtype,
        )
        bias_tensor = None
1205
        if self.use_bias:
1206
1207
1208
1209
1210
            bias_tensor = torch.empty(
                self.out_features,
                device=device,
                dtype=params_dtype,
            )
1211

1212
1213
1214
1215
        # Configure parameter splits
        self.weight_names = []
        self.bias_names = []
        self.parameter_split_sizes = []
1216
        if parameters_split is None:
1217
1218
1219
1220
1221
1222
            # Split into a single parameter by default
            self.weight_names = ["weight"]
            self.bias_names = ["bias"]
            self.parameter_split_sizes = [out_features]
        elif not parameters_split:
            raise ValueError("Cannot split weight buffer into 0 parameters")
cyanguwa's avatar
cyanguwa committed
1223
        elif isinstance(parameters_split, dict):
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
            # Split parameters with provided sizes
            for name, split_size in parameters_split.items():
                self.weight_names.append(f"{name.rstrip('_')}_weight")
                self.bias_names.append(f"{name.rstrip('_')}_bias")
                self.parameter_split_sizes.append(split_size)
        elif all(isinstance(name, str) for name in parameters_split):
            # Split parameters evenly
            split_size = out_features // len(parameters_split)
            for name in parameters_split:
                self.weight_names.append(f"{name.rstrip('_')}_weight")
                self.bias_names.append(f"{name.rstrip('_')}_bias")
                self.parameter_split_sizes.append(split_size)
cyanguwa's avatar
cyanguwa committed
1236
        else:
1237
            raise TypeError("Invalid configuration for parameters split")
1238

1239
1240
1241
1242
1243
1244
        # Make sure parameter splits are valid
        if sum(self.parameter_split_sizes) != out_features:
            raise ValueError(
                f"Trying to split weight buffer ({out_features=}) "
                f"with split sizes {self.parameter_split_sizes}"
            )
1245

1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
        # Adjust parameter splits for tensor-parallel distribution
        if self.parallel_mode == "column":
            for i, size in enumerate(self.parameter_split_sizes):
                if size % self.tp_size != 0:
                    raise RuntimeError(
                        f"Attempting to distribute a parameter with out_features={size} "
                        f"between {self.tp_size} tensor-parallel processes"
                    )
                self.parameter_split_sizes[i] = size // self.tp_size

1256
1257
1258
1259
1260
        # Construct weight parameters
        # Note: Register weights together so that they are adjacent to
        # each other in Linear.parameters(). This makes it more likely
        # that they will stay contiguous if the weights are
        # manipulated externally, e.g. by FSDP.
1261
1262
1263
1264
1265
1266
1267
1268
        offset = 0
        for i, split_size in enumerate(self.parameter_split_sizes):
            split_start = offset
            offset += split_size
            split_end = offset

            # Check if parameters are subviews of buffers
            is_subview = (split_start, split_end) != (0, self.out_features)
1269
            if is_subview and with_fp8_params:
1270
1271
1272
                raise RuntimeError(
                    "Splitting QuantizedTensor into multiple params is not supported"
                )
1273

1274
            # Construct weight parameter
1275
1276
1277
1278
1279
1280
1281
            self.register_parameter(
                self.weight_names[i],
                torch.nn.Parameter(weight_tensor[split_start:split_end]),
                init_fn=init_method,
                get_rng_state_tracker=get_rng_state_tracker,
                fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT,
            )
1282

1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
        # Construct bias parameters if needed
        if self.use_bias:
            offset = 0
            for i, split_size in enumerate(self.parameter_split_sizes):
                split_start = offset
                offset += split_size
                split_end = offset
                self.register_parameter(
                    self.bias_names[i],
                    torch.nn.Parameter(bias_tensor[split_start:split_end]),
                    init_fn=init_method_constant(0.0),
                )
        else:
            for name in self.bias_names:
                bias = torch.Tensor().to(dtype=params_dtype, device=device)
                setattr(self, name, bias)
cyanguwa's avatar
cyanguwa committed
1299

1300
        if with_fp8_params:
1301
1302
            self.init_fp8_metadata()

1303
        self.reset_parameters(defer_init=device == "meta")
1304

1305
1306
1307
1308
1309
1310
1311
        # For RPL, bias has to be added after TP collectives
        # So it cannot be fused with the GEMM
        if self.parallel_mode == "row" and self.apply_bias:
            self.gemm_bias_unfused_add = True
        else:
            self.gemm_bias_unfused_add = False

1312
1313
1314
1315
1316
        if self.wgrad_store.delay_wgrad_compute():
            for name, param in self.named_parameters():
                if name in self.weight_names or name in self.bias_names:
                    param.skip_backward_post_hook = True

1317
1318
1319
1320
    def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
        """Init scales and amaxes for fwd | bwd."""
        super().set_meta_tensor(fwd, recipe)

1321
        # Recipe-specific quantizer configuration
1322
1323
1324
        recipe = FP8GlobalStateManager.get_fp8_recipe()
        if recipe.float8_current_scaling():
            self._customize_quantizers_float8_current_scaling(fwd, recipe)
1325
1326
        elif recipe.nvfp4():
            self._customize_quantizers_nvfp4(fwd, recipe)
1327

1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
    def reset_parameters(self, defer_init=False):
        super().reset_parameters(defer_init=defer_init)

        if not defer_init:
            # Set parallelism attributes for linear weights
            for weight in self.weight_names:
                set_tensor_model_parallel_attributes(
                    tensor=getattr(self, weight),
                    is_parallel=True,
                    dim=1 if self.parallel_mode == "row" else 0,
                    stride=1,
                )

            # Set parallelism attributes for linear biases
            if self.use_bias:
                for bias in self.bias_names:
                    if self.parallel_mode == "row":
                        setattr(getattr(self, bias), "sequence_parallel", self.sequence_parallel)
                    elif self.parallel_mode == "column":
                        set_tensor_model_parallel_attributes(getattr(self, bias), True, 0, 1)

1349
    @no_torch_dynamo()
1350
1351
1352
1353
    def forward(
        self,
        inp: torch.Tensor,
        is_first_microbatch: Optional[bool] = None,
1354
        fp8_output: Optional[bool] = False,
1355
        fp8_grad: Optional[bool] = False,
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]:
        """
        Apply the linear transformation to the input.

        Parameters
        ----------
        inp : torch.Tensor
             Input tensor.
        is_first_microbatch : {True, False, None}, default = None
                             During training using either gradient accumulation or
                             pipeline parallelism a minibatch of data is further split
                             into microbatches. Between the microbatches of the same minibatch
                             the model weights are not updated. Setting this parameter indicates
                             whether the current microbatch is the first in a minibatch or not.
                             When set, this parameter enables additional optimizations:

                             * during FP8 training, it allows caching of the FP8 versions of
                               the weights
                             * it also allows skipping gradient accumulation during the
                               first microbatch (since it is the first gradient being
                               produced)
        """
1378
1379
        is_grad_enabled = torch.is_grad_enabled()

1380
        if is_in_onnx_export_mode():
1381
            return self.onnx_forward(inp, fp8_output, is_grad_enabled)
1382

1383
        debug = self.is_debug_iter()
1384

1385
1386
1387
1388
        if FP8GlobalStateManager.fp8_graph_capturing():
            skip_fp8_weight_update = FP8GlobalStateManager.get_skip_fp8_weight_update_tensor()
        else:
            skip_fp8_weight_update = None
1389
1390
1391
        if skip_fp8_weight_update is not None:
            is_first_microbatch = False

1392
        if self.ub_overlap_rs_fprop:
1393
1394
1395
            if get_ub(
                self.ub_name + "_fprop", FP8GlobalStateManager.is_fp8_enabled()
            ).is_fp8_ubuf():
1396
1397
                fp8_output = True
        if self.ub_overlap_rs_dgrad:
1398
1399
1400
            if get_ub(
                self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled()
            ).is_fp8_ubuf():
1401
1402
                fp8_grad = True

1403
        with self.prepare_forward(
1404
            inp,
1405
            allow_non_contiguous=isinstance(inp, QuantizedTensor),
1406
        ) as inp:
1407

1408
            weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
1409

1410
            quantizers = (
1411
                self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
1412
                if not debug
1413
                else self._get_debug_quantizers(fp8_output, fp8_grad, is_grad_enabled)
1414
1415
            )
            if debug:
1416
                if self.no_debug_features_active(quantizers):
1417
                    debug = False
1418
                    quantizers = self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
1419

1420
1421
1422
1423
1424
            (
                input_quantizer,
                weight_quantizer,
                output_quantizer,
                grad_input_quantizer,
1425
1426
1427
                grad_weight_quantizer,
                grad_output_quantizer,
            ) = quantizers
1428

1429
            if is_grad_enabled:
1430
                linear_fn = _Linear.apply
1431
                autograd_ctx = []
1432
1433
            else:
                linear_fn = _Linear.forward
1434
1435
1436
                autograd_ctx = [None]

            non_tensor_args = (
1437
1438
1439
                is_first_microbatch,
                self.fp8,
                self.fp8_calibration,
1440
                self.wgrad_store,
1441
1442
1443
1444
                input_quantizer,
                weight_quantizer,
                output_quantizer,
                grad_input_quantizer,
1445
1446
                grad_weight_quantizer,
                grad_output_quantizer,
1447
                self.fuse_wgrad_accumulation,
1448
                is_cpu_offload_enabled(),
1449
1450
1451
1452
1453
1454
                self.tp_group,
                self.tp_size,
                self.sequence_parallel,
                self.tp_size > 1,
                self.activation_dtype,
                self.parallel_mode,
1455
                is_grad_enabled,
1456
1457
1458
1459
1460
1461
                self.ub_overlap_rs_fprop,
                self.ub_overlap_ag_dgrad,
                self.ub_overlap_ag_fprop,
                self.ub_overlap_rs_dgrad,
                self.ub_bulk_dgrad,
                self.ub_bulk_wgrad,
1462
                self.ub_name,
1463
                fp8_output,
1464
                self.fsdp_group,
1465
1466
                self,
                skip_fp8_weight_update,
1467
                self.symmetric_ar_type,
1468
                self.save_original_input,
1469
                debug,
1470
            )
1471
1472
1473
1474
1475
1476
1477
            out = linear_fn(
                *autograd_ctx,
                weight_tensor,
                inp,
                bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None,
                non_tensor_args,
            )
1478
1479
1480
1481
1482
1483
        if self.gemm_bias_unfused_add:
            out = out + cast_if_needed(bias_tensor, self.activation_dtype)

        if self.return_bias:
            return out, cast_if_needed(bias_tensor, self.activation_dtype)
        return out
1484

1485
    def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled):
1486
        if not self.fp8:
1487
            return [None] * 6
1488
        grad_input_quantizer = None
1489
        grad_weight_quantizer = None
1490
1491
1492
        grad_output_quantizer = None
        output_quantizer = None
        input_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT]
1493
        input_quantizer.internal = True
1494
1495
        if not (self.parallel_mode == "column" and self.sequence_parallel):
            input_quantizer.optimize_for_gemm = True
1496
        (weight_quantizer,) = self._get_weight_quantizers()
1497
1498
        if fp8_output:
            output_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT]
1499
        if is_grad_enabled:
1500
1501
            grad_output_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1]
            grad_output_quantizer.internal = True
1502
1503
            if not (self.parallel_mode == "row" and self.sequence_parallel):
                grad_output_quantizer.optimize_for_gemm = True
1504
1505
1506
1507
1508
1509
1510
            if fp8_grad:
                grad_input_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1]
        return (
            input_quantizer,
            weight_quantizer,
            output_quantizer,
            grad_input_quantizer,
1511
1512
1513
1514
            grad_weight_quantizer,
            grad_output_quantizer,
        )

1515
1516
    def _get_debug_quantizers(self, fp8_output, fp8_grad, is_grad_enabled):
        original_quantizers = self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
1517
1518
1519
1520
1521
1522
1523
        assert TEDebugState.debug_enabled
        from ...debug.pytorch.debug_quantization import DebugQuantizer

        names = ["activation", "weight", "output", "dgrad", "wgrad", "gradient"]
        return tuple(
            DebugQuantizer(self.name, name, q, self.tp_group)
            for name, q in zip(names, original_quantizers)
1524
        )
1525

1526
    def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage]]:
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
        """Get the weight tensors of the module."""
        unfused_weights = [getattr(self, name) for name in self.weight_names]
        if any(isinstance(w, QuantizedTensor) for w in unfused_weights):
            if self.fp8:
                if len(unfused_weights) != 1:
                    raise RuntimeError(
                        "Splitting QuantizedTensor into multiple params is not supported"
                    )
            else:
                warnings.warn(
                    "You are using quantized weights without quantized compute. "
                    "Please make sure this is intentional."
                )
                unfused_weights = [w.dequantize() for w in unfused_weights]
        return unfused_weights

    def _get_weight_and_bias_tensors(self) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        # Get concatenated weight and bias tensors
        unfused_weights = self._get_weight_tensors()
        weight_tensor = noop_cat(unfused_weights)
        if self.use_bias:
            bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names])
        else:
            bias_tensor = None
        return weight_tensor, bias_tensor

    def onnx_forward(
        self,
        inp: torch.Tensor,
        fp8_output: bool,
1557
        is_grad_enabled: bool,
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
    ) -> torch.Tensor:
        """
        ONNX-compatible version of the forward function that provides numerical equivalence
        while only using operations that have defined ONNX symbolic translations.
        This simplified implementation is designed specifically for inference scenarios.
        """
        from ..export import onnx_gemm

        assert_warmed_up(self)
        assert not TEDebugState.debug_enabled, "Debug mode is not supported in ONNX export."
        weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
        (
            input_quantizer,
            weight_quantizer,
            output_quantizer,
            *_,
1574
        ) = self._get_quantizers(fp8_output, False, is_grad_enabled)
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
        inp_dtype = inp.dtype

        if input_quantizer is not None:
            inp_q = input_quantizer.onnx_quantize(inp)
            inp = input_quantizer.onnx_dequantize(inp_q)
            inp = inp.to(inp_dtype)

        if weight_quantizer is not None:
            weight_q = weight_quantizer.onnx_quantize(weight_tensor)
            weight_tensor = weight_quantizer.onnx_dequantize(weight_q)
        if bias_tensor is not None:
            bias_tensor = bias_tensor.to(inp_dtype)
        weight_tensor = weight_tensor.to(inp_dtype)

        if self.apply_bias:
            output = onnx_gemm(weight_tensor, inp, bias_tensor)
        else:
            output = onnx_gemm(weight_tensor, inp, None)

        if output_quantizer is not None:
            raise NotImplementedError("ONNX export of quantized output is not supported")

        if self.return_bias:
            return output, bias_tensor

        return output

1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
    def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe) -> None:
        """Customize quantizers based on current scaling recipe + linear."""
        assert (
            recipe.float8_current_scaling()
        ), "current scaling recipe quantizer customization here"
        if fwd:
            # set configs about amax epsilon and power_2_scale
            self.quantizers["scaling_fwd"][
                tex.FP8FwdTensors.GEMM1_INPUT
            ].force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale
            self.quantizers["scaling_fwd"][
                tex.FP8FwdTensors.GEMM1_INPUT
            ].amax_epsilon = recipe.fp8_quant_fwd_inp.amax_epsilon
            # also set weight quantizer with same amax_epsilon & power_2_scale
            self.quantizers["scaling_fwd"][
                tex.FP8FwdTensors.GEMM1_WEIGHT
            ].force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale
            self.quantizers["scaling_fwd"][
                tex.FP8FwdTensors.GEMM1_WEIGHT
            ].amax_epsilon = recipe.fp8_quant_fwd_weight.amax_epsilon
            # paralle related
            if self.sequence_parallel and self.parallel_mode == "column":
                # customize input_quantizer with amax reduction TP group
                self.quantizers["scaling_fwd"][
                    tex.FP8FwdTensors.GEMM1_INPUT
                ].with_amax_reduction = True
                self.quantizers["scaling_fwd"][
                    tex.FP8FwdTensors.GEMM1_INPUT
                ].amax_reduction_group = self.tp_group
        else:
            # set grad_output_quantizer with amax epsilon and power_2_scale
            self.quantizers["scaling_bwd"][
                tex.FP8BwdTensors.GRAD_OUTPUT1
            ].force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale
            self.quantizers["scaling_bwd"][
                tex.FP8BwdTensors.GRAD_OUTPUT1
            ].amax_epsilon = recipe.fp8_quant_bwd_grad.amax_epsilon
            # parallel related
            if self.sequence_parallel and self.parallel_mode == "row":
                # customize grad_output_quantizer with amax reduction TP group
                self.quantizers["scaling_bwd"][
                    tex.FP8BwdTensors.GRAD_OUTPUT1
                ].with_amax_reduction = True
                self.quantizers["scaling_bwd"][
                    tex.FP8BwdTensors.GRAD_OUTPUT1
                ].amax_reduction_group = self.tp_group
1648

1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
    def _customize_quantizers_nvfp4(self, fwd: bool, recipe: Recipe) -> None:
        """Customize quantizers based on current scaling recipe + linear."""
        assert recipe.nvfp4(), "Incorrect recipe."
        if fwd:
            if self.sequence_parallel and self.parallel_mode == "column":
                # customize input_quantizer with amax reduction TP group
                self.quantizers["scaling_fwd"][
                    tex.FP8FwdTensors.GEMM1_INPUT
                ].with_amax_reduction = True
                self.quantizers["scaling_fwd"][
                    tex.FP8FwdTensors.GEMM1_INPUT
                ].amax_reduction_group = self.tp_group
        else:
            if self.sequence_parallel and self.parallel_mode == "row":
                # customize grad_output_quantizer with amax reduction TP group
                self.quantizers["scaling_bwd"][
                    tex.FP8BwdTensors.GRAD_OUTPUT1
                ].with_amax_reduction = True
                self.quantizers["scaling_bwd"][
                    tex.FP8BwdTensors.GRAD_OUTPUT1
                ].amax_reduction_group = self.tp_group

1671
1672
    def _get_weight_quantizers(self) -> List[Quantizer]:
        """Get the weight quantizers of the module."""
1673
        if not self.fp8 and not self.fp8_calibration:
1674
1675
1676
1677
            return [None]
        weight_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT]
        weight_quantizer.internal = True
        return [weight_quantizer]