linear.py 57.4 KB
Newer Older
1
# Copyright (c) 2022-2025, 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
7
8
from functools import reduce
from operator import mul as multiply_op
9
10
11

import torch

12
import transformer_engine_torch as tex
13

14
from transformer_engine.common.recipe import Recipe
15
16
17
18
from .base import (
    get_workspace,
    get_ub,
    TransformerEngineBaseModule,
19
    get_dummy_wgrad,
20
21
22
23
    _2X_ACC_FPROP,
    _2X_ACC_DGRAD,
    _2X_ACC_WGRAD,
)
24
25
from ._common import noop_cat, _fix_gathered_fp8_transpose
from ..fp8 import FP8GlobalStateManager
26
27
from ..utils import (
    cast_if_needed,
28
    clear_tensor_data,
29
    divide,
30
    init_method_constant,
31
32
    requires_grad,
    needs_quantized_gemm,
33
    non_tn_fp8_gemm_supported,
34
    assert_dim_for_fp8_exec,
35
36
    nvtx_range_pop,
    nvtx_range_push,
37
38
39
40
41
42
43
)
from ..distributed import (
    set_tensor_model_parallel_attributes,
    get_distributed_world_size,
    allreduce,
    reduce_scatter_along_first_dim,
    gather_along_first_dim,
44
    is_fp8_activation_recompute_enabled,
45
    in_fp8_activation_recompute_phase,
46
47
    _fsdp_scatter_tensors,
    _fsdp_gather_tensors,
48
49
)
from ..cpp_extensions import (
50
    general_gemm,
51
)
52
from ..constants import GemmParallelModes, dist_group_type
53
from ..jit import no_torch_dynamo
54
from ..graph import is_graph_capturing
55
56
57
58
59
60
from ..tensor.quantized_tensor import (
    QuantizedTensor,
    Quantizer,
    prepare_for_saving,
    restore_from_saved,
)
61
from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer
62
from ..tensor.mxfp8_tensor import MXFP8Quantizer
63
from ..tensor._internal.mxfp8_tensor_base import MXFP8TensorBase
64
from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer
65
from ..cpu_offload import is_cpu_offload_enabled, set_offloading_param
66
67
from ...debug.pytorch.debug_state import TEDebugState
from ...debug.pytorch.utils import any_feature_enabled
68

69

70
71
72
73
74
75
76
77
78
79
80
__all__ = ["Linear"]


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

    @staticmethod
    def forward(
        ctx,
81
        weight: torch.Tensor,
82
        inp: torch.Tensor,
83
        bias: Optional[torch.Tensor],
84
85
86
        is_first_microbatch: Union[bool, None],
        fp8: bool,
        fp8_calibration: bool,
87
88
89
90
        input_quantizer: Optional[Quantizer],
        weight_quantizer: Optional[Quantizer],
        output_quantizer: Optional[Quantizer],
        grad_input_quantizer: Optional[Quantizer],
91
92
        grad_weight_quantizer: Optional[Quantizer],
        grad_output_quantizer: Optional[Quantizer],
93
        fuse_wgrad_accumulation: bool,
94
        cpu_offloading: bool,
95
96
97
98
99
100
101
        tp_group: Union[dist_group_type, None],
        tp_size: int,
        sequence_parallel: bool,
        tensor_parallel: bool,
        activation_dtype: torch.dtype,
        parallel_mode: Union[str, None],
        is_grad_enabled: bool,
102
103
104
105
106
107
        ub_overlap_rs_fprop: bool,
        ub_overlap_ag_dgrad: bool,
        ub_overlap_ag_fprop: bool,
        ub_overlap_rs_dgrad: bool,
        ub_bulk_dgrad: bool,
        ub_bulk_wgrad: bool,
108
        ub_name: str,
109
        fp8_output: bool,  # pylint: disable=unused-argument
110
        fsdp_group: Union[dist_group_type, None],
111
112
        module: torch.nn.Module,
        skip_fp8_weight_update: bool,
113
        debug: Optional[bool] = False,
114
    ) -> torch.Tensor:
115
        # pylint: disable=missing-function-docstring
116

117
118
119
120
121
        # NVTX label for profiling
        nvtx_label = "transformer_engine._Linear.forward"
        if ub_name is not None:
            nvtx_label = f"{nvtx_label}.{ub_name}"

122
        # Make sure input dimensions are compatible
123
124
125
        out_features, in_features = weight.shape
        inp_shape = inp.shape
        assert inp_shape[-1] == in_features, "GEMM not possible"
126

127
        tp_world_size = get_distributed_world_size(tp_group)
128
129
130
131
        backward_needs_input = is_grad_enabled and weight.requires_grad

        # Prepare input tensor
        # Note: Cast to expected dtype and perform tensor-parallel communication
132
        nvtx_range_push(f"{nvtx_label}.input_cast_comm")
133
        inputmat = inp.view(-1, in_features)
134
135
136
137
138
        inputmat_total = None
        with_input_all_gather_nccl = (
            parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop
        )
        own_quantized_input = False
139
140
141
142
        # TODO(kwyss): Support FP8 allgather for FP8 block quantization.
        force_hp_input_gather = (
            fp8 and with_input_all_gather_nccl and isinstance(input_quantizer, Float8BlockQuantizer)
        )  # Perform TP communication in high precision.
143
        if fp8:
144
            assert_dim_for_fp8_exec(inputmat, weight)
145
146
            if any([ub_overlap_ag_fprop, ub_overlap_rs_fprop]) and not (
                FP8GlobalStateManager.get_fp8_recipe().float8_per_tensor_scaling()
147
148
            ):
                raise NotImplementedError(
149
150
                    "Comm+GEMM overlap is only supported with FP8 delayed scaling or per-tensor"
                    " current scaling"
151
                )
152
        if fp8 or debug:
153
154
155
            if input_quantizer is None:
                raise ValueError("Missing quantizer for input tensor")
            if with_input_all_gather_nccl:
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
                if force_hp_input_gather:
                    input_quantizer.set_usage(rowwise=True, columnwise=False)
                    inputmat_total, _ = gather_along_first_dim(
                        inputmat, tp_group, quantizer=input_quantizer
                    )
                else:
                    if not isinstance(inputmat, QuantizedTensor):
                        columnwise_usage = backward_needs_input and isinstance(
                            input_quantizer, MXFP8Quantizer
                        )
                        # force_hp_input_gather should enforce this
                        assert not isinstance(input_quantizer, Float8BlockQuantizer)
                        input_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)
                        inputmat = input_quantizer(inputmat)
                        own_quantized_input = True
                    input_quantizer.set_usage(rowwise=True, columnwise=False)
                    inputmat_total, _ = gather_along_first_dim(
                        inputmat,
                        tp_group,
                        quantizer=input_quantizer,
176
                    )
177
            else:
178
179
180
181
182
183
184
185
186
187
188
                if (
                    FP8GlobalStateManager.get_fp8_recipe().float8_per_tensor_scaling()
                    and ub_bulk_dgrad
                ):
                    # reduce duplicated transpose in `_fix_gathered_fp8_transpose`
                    input_quantizer.set_usage(rowwise=True, columnwise=False)
                else:
                    input_quantizer.set_usage(
                        rowwise=True,
                        columnwise=backward_needs_input,
                    )
189
190
                if not isinstance(inputmat, QuantizedTensor):
                    inputmat = input_quantizer(inputmat)
191
                    own_quantized_input = True
192
193
194
                elif backward_needs_input:
                    inputmat.update_usage(rowwise_usage=True, columnwise_usage=True)
                inputmat_total = inputmat
195
        else:
196
197
198
            inputmat = cast_if_needed(inp, activation_dtype)
            if with_input_all_gather_nccl:
                inputmat_total, _ = gather_along_first_dim(inputmat, tp_group)
199
            else:
200
                inputmat_total = inputmat
201
        nvtx_range_pop(f"{nvtx_label}.input_cast_comm")
202

203
        # Cast weight to expected dtype
204
205
206
        weightmat = weight

        if fp8 or debug:
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
            # Configure quantizer
            if weight_quantizer is not None:
                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)
            # FP8 cast to workspace buffer
            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,
225
                workspace_dtype=activation_dtype,
226
            )
227
228
        else:
            weightmat = cast_if_needed(weightmat, activation_dtype)
229
230
231

        # Cast bias to expected dtype
        bias_dtype = activation_dtype
232
        if needs_quantized_gemm(inputmat_total) and activation_dtype == torch.float32:
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
            bias_dtype = torch.bfloat16
        bias = cast_if_needed(bias, bias_dtype) if bias is not None else bias

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

        # 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)

        ub_obj = None
        ub_type = None
        rs_out = None
        out_dtype = activation_dtype
        if ub_overlap_rs_fprop:
            ub_obj = get_ub(ub_name + "_fprop")
            ub_type = tex.CommOverlapType.RS
            out_shape = [reduce(multiply_op, inp_shape[:-1]) // tp_world_size, out_features]
            rs_out = torch.empty(out_shape, dtype=activation_dtype, device=inputmat_total.device)

        elif ub_overlap_ag_fprop:
            ub_obj = get_ub(ub_name + "_fprop")
            ub_type = tex.CommOverlapType.AG
            if fp8:
                assert ub_obj.is_fp8_ubuf(), "AG overlap with FP8 GEMM inputs requires FP8 buffer."
            ub_obj.copy_into_buffer(inputmat_total, input_quantizer, local_chunk=True)
            inputmat_total = ub_obj.get_buffer(input_quantizer)

265
        nvtx_range_push(f"{nvtx_label}.gemm")
266
267
268
269
270
271
        fprop_gemm_use_split_accumulator = _2X_ACC_FPROP
        if fp8:
            recipe = FP8GlobalStateManager.get_fp8_recipe()
            if hasattr(recipe, "fp8_gemm_fprop"):
                fprop_gemm_use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator

272
273
274
275
276
277
278
        out, *_, rs_out = general_gemm(
            weightmat,
            inputmat_total,
            get_workspace(),
            quantization_params=output_quantizer,
            out_dtype=out_dtype,
            bias=bias,
279
            use_split_accumulator=fprop_gemm_use_split_accumulator,
280
281
282
283
            ub=ub_obj,
            ub_type=ub_type,
            extra_output=rs_out,
        )
284
        nvtx_range_pop(f"{nvtx_label}.gemm")
285
286

        if is_grad_enabled:
287
            ctx.weight_quantizer = weight_quantizer
288
            saved_inputmat = None
289
290
291
292
293

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

294
295
            if backward_needs_input:
                if own_quantized_input and isinstance(inputmat, QuantizedTensor):
296
297
298
299
                    # For sequence parallel in vanilla FP8, rowwise data is
                    # to gather the input. For MXFP8, columnwise only data
                    # can be allgathered.
                    if isinstance(inputmat, MXFP8TensorBase) or not ctx.backward_input_needs_gather:
300
                        inputmat.update_usage(rowwise_usage=False, columnwise_usage=True)
301
302
                if force_hp_input_gather:
                    assert not isinstance(inputmat, QuantizedTensor)
303
                saved_inputmat = inputmat
304

305
306
307
308
309
            # Weight with column-wise usage is needed for dgrad GEMM.
            if inp.requires_grad:
                if isinstance(weightmat, QuantizedTensor):
                    weightmat.update_usage(columnwise_usage=True)

310
311
312
313
314
            if cpu_offloading:
                set_offloading_param(weight, "weight_offloading", True)
                set_offloading_param(weightmat, "weight_offloading", True)
                if saved_inputmat is not None:
                    set_offloading_param(saved_inputmat, "activation_offloading", True)
315

316
317
            # Scatter intermediate/activation tensors saved for the backward pass
            # NOTE: FSDP sharding is not valid for models initialized with primary Fp8 weights
318
            nvtx_range_push(f"{nvtx_label}.fsdp_scatter")
319
320
321
            ctx.fsdp_group = fsdp_group
            ctx.fsdp_shapes = _fsdp_scatter_tensors(
                fsdp_group,
322
323
                saved_inputmat,
                weightmat if fp8 and not isinstance(weight, QuantizedTensor) else None,
324
            )
325
            nvtx_range_pop(f"{nvtx_label}.fsdp_scatter")
326

327
328
329
330
331
332
333
334
335
336
337
            if cpu_offloading:
                ctx.grad_added_to_main_grad = hasattr(weight, "grad_added_to_main_grad")

                if ctx.grad_added_to_main_grad:
                    # 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

338
339
            # TODO(ksivamani): Check memory usage
            tensors_to_save, tensor_objects = prepare_for_saving(
340
                saved_inputmat,
341
                weightmat,
342
                weight,
343
                bias,
344
            )
345
346
            ctx.save_for_backward(*tensors_to_save)
            ctx.tensor_objects = tensor_objects
347

348
349
            ctx.activation_dtype = activation_dtype
            ctx.fp8 = fp8
350
351
            ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None
            ctx.force_hp_input_gather = force_hp_input_gather
352
353
            ctx.input_quantizer = input_quantizer
            ctx.grad_input_quantizer = grad_input_quantizer
354
355
            ctx.grad_weight_quantizer = grad_weight_quantizer
            ctx.grad_output_quantizer = grad_output_quantizer
356
            ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation
357
358
359
            if fuse_wgrad_accumulation and weight.requires_grad:
                ctx.main_grad = weight.main_grad

360
            ctx.debug = debug
361
            ctx.cpu_offloading = cpu_offloading
362
            ctx.is_first_microbatch = is_first_microbatch
363
            ctx.use_bias = bias is not None
364
365
            ctx.sequence_parallel = sequence_parallel
            ctx.tensor_parallel = tensor_parallel
366
            ctx.inp_shape = inp_shape
367
368
            ctx.parallel_mode = parallel_mode
            ctx.tp_group = tp_group
369
370
371
372
            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
373
            ctx.ub_name = ub_name
374
375
            ctx.tp_size = tp_size
            ctx.requires_dgrad = inp.requires_grad
376
            ctx.requires_wgrad = weight.requires_grad
377
            ctx.reduce_and_update_bwd_fp8_tensors = False
378
            ctx.owns_input = saved_inputmat is not inp
379
            if ctx.fp8 and requires_grad(inp, weight, bias):
380
381
382
383
                _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
384
385

        # Row Parallel Linear
386
387
388
        if ub_overlap_rs_fprop:
            out = rs_out
        elif parallel_mode == "row":
389
            nvtx_range_push(f"{nvtx_label}.row_parallel_comm")
390
            if sequence_parallel:
391
392
393
                out, _ = reduce_scatter_along_first_dim(out, tp_group)
            elif tensor_parallel:
                out, _ = allreduce(out, tp_group)
394
            nvtx_range_pop(f"{nvtx_label}.row_parallel_comm")
395

396
397
        out = out.view(-1, *inp_shape[1:-1], out_features)
        return out
398
399

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

403
404
405
406
407
        # 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}"

408
        with torch.cuda.nvtx.range("_Linear_backward"):
409
410
411
412
413
414
415
416
417
418
            if (
                ctx.fp8
                and any(
                    [
                        ctx.ub_overlap_ag,
                        ctx.ub_overlap_rs_dgrad,
                        ctx.ub_bulk_dgrad,
                        ctx.ub_bulk_wgrad,
                    ]
                )
419
                and (ctx.fp8_recipe is not None)
420
            ):
421
                if not ctx.fp8_recipe.float8_per_tensor_scaling():
422
                    raise NotImplementedError(
423
424
                        "Comm+GEMM overlap is only supported with FP8 delayed scaling or per-tensor"
                        " current scaling"
425
                    )
426
427
428
429
430

            saved_tensors = ctx.saved_tensors
            inputmat, weight_fp8, weight, bias = (  # pylint: disable=unbalanced-tuple-unpacking
                restore_from_saved(ctx.tensor_objects, saved_tensors)
            )
431
432
433
            # 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
434
435
436
437
438
439
440
441

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

442
443
444
445
446
            if ctx.cpu_offloading:
                if ctx.grad_added_to_main_grad:
                    weight = ctx.weight_object
                if ctx.requires_wgrad and ctx.fuse_wgrad_accumulation:
                    weight.main_grad = main_grad
447

448
449
450
            # 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
451
            nvtx_range_push(f"{nvtx_label}.fsdp_gather")
452
453
454
455
            _fsdp_gather_tensors(
                ctx.fsdp_group,
                ctx.fsdp_shapes,
                inputmat,
456
                weight_fp8,
457
            )
458
            nvtx_range_pop(f"{nvtx_label}.fsdp_gather")
459

460
            ctx.ub_obj_gradout = None
461
            ub_obj_dgrad = None
462
            ub_obj_wgrad = None
463
464
            ub_type_dgrad = None
            ub_type_wgrad = None
465
            dgrad_shape = [reduce(multiply_op, ctx.inp_shape[:-1]), ctx.inp_shape[-1]]
466
467
            rs_out = None
            dgrad_bulk = None
468
            if ctx.ub_overlap_ag:
469
                # Overlap grad_output all-gather with dgrad compute
470
                ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad")
471
472
                ub_obj_dgrad = ctx.ub_obj_gradout
                ub_type_dgrad = tex.CommOverlapType.AG
473
474
475
476

            elif ctx.ub_overlap_rs_dgrad:
                # Overlap dgrad reduce-scatter with dgrad compute
                ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad")
477
478
                ub_obj_dgrad = ctx.ub_obj_gradout
                ub_type_dgrad = tex.CommOverlapType.RS
479
480
481
482
483
484
485
                rs_out = torch.empty(
                    dgrad_shape, dtype=ctx.activation_dtype, device=grad_output.device
                )

            else:
                if ctx.ub_bulk_dgrad:
                    # Overlap inputmat all-gather with dgrad compute
486
487
488
489
                    # NOTE: Copying into communication buffer will always prefer rowwise data,
                    #       and will copy columnwise data if rowwise does not exist. In that case,
                    #       the all-gather will apply to the leading dimension of the transpose,
                    #       which then needs to be interleaved correctly before WGRAD.
490
                    ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad")
491
492
493
                    ub_obj_dgrad = ctx.ub_obj_gradout
                    ub_type_dgrad = tex.CommOverlapType.AG
                    ub_obj_dgrad.copy_into_buffer(inputmat, ctx.input_quantizer, local_chunk=True)
494
495
496
497

                if ctx.ub_bulk_wgrad:
                    # Overlap dgrad reduce-scatter with wgrad compute
                    ub_obj_wgrad = get_ub(ctx.ub_name + "_wgrad")
498
499
500
501
                    ub_type_wgrad = tex.CommOverlapType.RS
                    ub_obj_wgrad.set_buffer_params(ctx.grad_input_quantizer)
                    dgrad_bulk = ub_obj_wgrad.get_buffer(ctx.grad_input_quantizer)

502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
            # 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:
                rowwise_usage = True
                columnwise_usage = True
                if ctx.ub_overlap_ag and isinstance(
                    ctx.grad_output_quantizer,
                    (Float8Quantizer, Float8CurrentScalingQuantizer),
                ):
                    # If data is in FP8 and communication is handled
                    # with Userbuffers, we compute FP8 transposes
                    # manually
                    columnwise_usage = False
                ctx.grad_output_quantizer.set_usage(
                    rowwise=rowwise_usage,
                    columnwise=columnwise_usage,
                )

521
522
            # Prepare grad output tensor
            # Note: Cast to expected dtype and perform tensor-parallel communication
523
            nvtx_range_push(f"{nvtx_label}.grad_output_preprocess")
524
525
526
527
            (
                grad_output,
                grad_bias,
            ) = TransformerEngineBaseModule.grad_output_preprocess(
528
529
530
531
                ctx,
                grad_output,
                ctx.parallel_mode == "row",
                ctx.grad_output_quantizer,
532
            )
533
            nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess")
534

535
            # Launch tensor-parallel communication for input tensor
536
            inputmat_total = None
537
            inputmat_total_work = None
538
            if ctx.backward_input_needs_gather and not ctx.ub_bulk_dgrad:
539
                quantizer = None
540
                if ctx.fp8 or ctx.debug:
541
                    quantizer = ctx.input_quantizer
542
543
544
545
546
547
                    if isinstance(quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer)):
                        # 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)
548
                nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input")
549
                gather_quantizer = None if ctx.force_hp_input_gather else quantizer
550
551
552
553
                inputmat_total, inputmat_total_work = gather_along_first_dim(
                    inputmat,
                    ctx.tp_group,
                    async_op=True,
554
                    quantizer=gather_quantizer,
555
                )
556
                nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input")
557
558
559
            else:
                inputmat_total = inputmat

560
            # Check whether to output wgrad GEMM directly into main grad
561
562
563
564
565
566
567
            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

568
569
570
            # Compute grad input tensor
            dgrad = None
            dgrad_work = None
571
            if ctx.requires_dgrad:
572
573
574
575
576

                # Update quantizer
                if ctx.grad_input_quantizer is not None:
                    ctx.grad_input_quantizer.set_usage(rowwise=True, columnwise=False)
                # dgrad GEMM
577
                nvtx_range_push(f"{nvtx_label}.dgrad_gemm")
578
579
580
581
582
583
584
585
                dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD
                if ctx.fp8:
                    recipe = ctx.fp8_recipe
                    if hasattr(recipe, "fp8_gemm_dgrad"):
                        dgrad_gemm_use_split_accumulator = (
                            recipe.fp8_gemm_dgrad.use_split_accumulator
                        )

586
587
588
589
590
591
                if ctx.weight_quantizer is not None and isinstance(weight_fp8, QuantizedTensor):
                    weight_fp8.update_usage(
                        rowwise_usage=ctx.weight_quantizer.rowwise_usage,
                        columnwise_usage=ctx.weight_quantizer.columnwise_usage,
                    )

592
593
594
595
596
597
598
599
600
                dgrad, *_, rs_out = general_gemm(
                    weight_fp8,
                    grad_output,
                    get_workspace(),
                    layout="NN",
                    grad=True,
                    quantization_params=ctx.grad_input_quantizer,
                    out=dgrad_bulk,
                    out_dtype=ctx.activation_dtype,
601
                    use_split_accumulator=dgrad_gemm_use_split_accumulator,
602
603
604
605
606
                    ub=ub_obj_dgrad,
                    ub_type=ub_type_dgrad,
                    extra_output=rs_out,
                    bulk_overlap=ctx.ub_bulk_dgrad,
                )
607
                nvtx_range_pop(f"{nvtx_label}.dgrad_gemm")
608
609
610
611
612

                # Launch tensor-parallel communication
                if ctx.ub_overlap_rs_dgrad:
                    dgrad = rs_out
                elif ctx.parallel_mode == "column" and not ctx.ub_bulk_wgrad:
613
                    nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad")
614
615
616
617
618
                    if ctx.sequence_parallel:
                        dgrad, dgrad_work = reduce_scatter_along_first_dim(
                            dgrad,
                            ctx.tp_group,
                            async_op=True,
619
                        )
620
                    else:
621
                        dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True)
622
                    nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad")
623

624
625
626
            # Compute grad weight tensor
            wgrad = None
            if ctx.requires_wgrad:
627
628

                # Synchronize tensor-parallel communication for input tensor
629
630
631
632
633
634
635
636
637
638
639
640
641
                if ctx.ub_bulk_dgrad:
                    inputmat_total = ub_obj_dgrad.get_buffer(ctx.input_quantizer)
                    if ctx.fp8:
                        if inputmat._data is None:
                            # All-gather executed on columnwise data and result is in rowwise data,
                            # so we need to fix the interleaving before WGRAD.
                            inputmat_total = _fix_gathered_fp8_transpose(
                                inputmat_total, ctx.tp_size
                            )
                        elif not non_tn_fp8_gemm_supported():
                            # FP8 GEMM on Hopper only supports TN layout so the gathered input must
                            # have a valid transpose.
                            inputmat_total._create_transpose()
642
643
644
                if inputmat_total_work is not None:
                    inputmat_total_work.wait()
                    inputmat_total_work = None
645
646
647
648
649
650
651
                if ctx.input_quantizer is not None and not isinstance(
                    inputmat_total, QuantizedTensor
                ):
                    # Async gather in BF16 does not asynchronously
                    # call quantizer after gather.
                    ctx.input_quantizer.set_usage(rowwise=False, columnwise=True)
                    inputmat_total = ctx.input_quantizer(inputmat_total)
652

653
654
655
                # Make sure GEMM inputs have required data
                if isinstance(inputmat_total, QuantizedTensor):
                    inputmat_total.update_usage(columnwise_usage=True)
656
                if isinstance(grad_output, QuantizedTensor):
657
                    grad_output.update_usage(columnwise_usage=True)
658

659
660
661
662
663
664
665
666
667
                # 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

                # Output buffer for overlapping grad input
                # reduce-scatter with wgrad GEMM
668
669
670
                if ctx.ub_bulk_wgrad and ub_obj_wgrad.is_fp8_ubuf():
                    rs_out = torch.empty(
                        dgrad_shape, dtype=ctx.activation_dtype, device=grad_output.device
671
672
                    )

673
674
                # wgrad GEMM
                # Note: Fuse with bgrad computation if needed
675
                nvtx_range_push(f"{nvtx_label}.wgrad_gemm")
676
677
678
679
680
681
682
683
684
685
686
                wgrad, grad_bias_, _, rs_out = general_gemm(
                    inputmat_total,
                    grad_output,
                    get_workspace(),
                    layout="NT",
                    grad=True,
                    out_dtype=(
                        main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype
                    ),
                    bias=(bias if (grad_bias is None and not ctx.fp8) else None),
                    out=main_grad if ctx.fuse_wgrad_accumulation else None,
687
                    use_split_accumulator=use_split_accumulator,
688
                    accumulate=accumulate_wgrad_into_param_main_grad,
689
                    quantization_params=ctx.grad_weight_quantizer,
690
691
692
693
694
                    ub=ub_obj_wgrad,
                    ub_type=ub_type_wgrad,
                    extra_output=rs_out,
                    bulk_overlap=ctx.ub_bulk_wgrad,
                )
695
                nvtx_range_pop(f"{nvtx_label}.wgrad_gemm")
696

697
698
699
                if ctx.ub_bulk_wgrad:
                    if ub_obj_wgrad.is_fp8_ubuf():
                        dgrad = rs_out
700
                    else:
701
                        dgrad = ub_obj_wgrad.get_buffer(ctx.grad_input_quantizer, local_chunk=True)
702

703
704
705
                if grad_bias is None:
                    grad_bias = grad_bias_
                del grad_bias_
706

707
                # Deallocate input tensor
708
709
                if ctx.owns_input:
                    clear_tensor_data(inputmat_total)
710

711
            # Don't return grad bias if not needed
712
713
714
            if not ctx.use_bias:
                grad_bias = None

715
            # Make sure all tensor-parallel communication is finished
716
717
718
719
720
721
722
723
            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:
724
            # Handle custom DDP from mcore.
725
726
727
728
729
            if (
                ctx.fuse_wgrad_accumulation
                and weight is not None
                and hasattr(weight, "grad_added_to_main_grad")
            ):
730
                weight.grad_added_to_main_grad = True
731
                if getattr(weight, "zero_out_wgrad", False):
732
733
734
735
                    wgrad = get_dummy_wgrad(
                        list(weight.main_grad.shape),
                        weight.dtype,
                        zero=True,
736
                    )
737
                else:
738
739
740
                    wgrad = get_dummy_wgrad(
                        list(weight.main_grad.shape),
                        weight.dtype,
741
                    )
742
743
744
745
            elif ctx.fuse_wgrad_accumulation:
                wgrad = None
        else:
            wgrad = None
746

747
        if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing():
748
            nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors")
749
            FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
750
            nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors")
751

752
        # Scatter fp8 weight buffers
753
        if ctx.fp8 and not isinstance(weight, QuantizedTensor):
754
            _fsdp_scatter_tensors(ctx.fsdp_group, weight_fp8)
755
        return (
756
            wgrad,
757
758
            dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None,
            grad_bias,
759
760
761
            None,  # is_first_microbatch
            None,  # fp8
            None,  # fp8_calibration
762
763
764
765
            None,  # input_quantizer
            None,  # weight_quantizer
            None,  # output_quantizer
            None,  # grad_input_quantizer
766
767
            None,  # grad_weight_quantizer
            None,  # grad_output_quantizer
768
769
770
771
772
773
774
775
776
            None,  # fuse_wgrad_accumulation
            None,  # cpu_offloading
            None,  # tp_group
            None,  # tp_size
            None,  # sequence_parallel
            None,  # tensor_parallel
            None,  # activation_dtype
            None,  # parallel_mode
            None,  # is_grad_enabled
777
778
779
780
781
782
            None,  # ub_overlap_rs_fprop
            None,  # ub_overlap_ag_dgrad
            None,  # ub_overlap_ag_fprop
            None,  # ub_overlap_rs_dgrad
            None,  # ub_bulk_dgrad
            None,  # ub_bulk_wgrad
783
            None,  # ub_name
784
            None,  # fp8_output
785
            None,  # fsdp_group
786
787
            None,  # module
            None,  # skip_fp8_weight_update
788
            None,  # debug
789
790
791
792
        )


class Linear(TransformerEngineBaseModule):
793
    """Applies a linear transformation to the incoming data :math:`y = xA^T + b`
794
795
796
797
798
799
800
801
802
803
804
805
806
807

    On NVIDIA GPUs it is a drop-in replacement for `torch.nn.Linear`.

    Parameters
    ----------
    in_features : int
                 size of each input sample.
    out_features : int
                  size of each output sample.
    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)`.
808
    get_rng_state_tracker : Callable, default = `None`
809
                 used to get the random number generator state tracker for initializing weights.
810
811
    rng_tracker_name : str, default = `None`
                 the param passed to get_rng_state_tracker to get the specific rng tracker.
cyanguwa's avatar
cyanguwa committed
812
    parameters_split : Optional[Union[Tuple[str, ...], Dict[str, int]]], default = None
813
814
815
816
817
818
819
                      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
                      names that end in `_weight` or `_bias`, so trailing underscores are
                      stripped from any provided names.
820
    device : Union[torch.device, str], default = "cuda"
821
          The device on which the parameters of the model will be allocated. It is the user's
822
823
          responsibility to ensure all parameters are moved to the GPU before running the
          forward pass.
824
825
    name: str, default = `None`
        name of the module, currently used for debugging purposes.
826
827
828
829
830
831
832
833
834
835
836
837
838

    Parallelism parameters
    ----------------------
    sequence_parallel : bool, default = `False`
                       if set to `True`, uses sequence parallelism.
    tp_group : ProcessGroup, default = `None`
              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
             `set_tensor_parallel_group(tp_group)` method on the initialized module before the
             forward pass to supply the tensor parallel group needed for tensor and sequence
             parallel collectives.
839
    parallel_mode : {None, 'column', 'row'}, default = `None`
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
                   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>`_.
                   When set to `None`, no communication is performed.

    Optimization parameters
    -----------------------
    fuse_wgrad_accumulation : bool, default = 'False'
                             if set to `True`, enables fusing of creation and accumulation of
                             the weight gradient. When enabled, it is assumed that the weights
                             have an additional `main_grad` attribute (used instead of the
                             regular `grad`) which is a pre-allocated buffer of the correct
                             size to accumulate gradients in.
    return_bias : bool, default = `False`
                 when set to `True`, this module will not apply the additive bias itself, but
                 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.
857
    params_dtype : torch.dtype, default = `torch.get_default_dtype()`
858
859
860
                  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.
861

862
863
864
865
866
867
868
869
870
871
872
    """

    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,
873
        rng_tracker_name: Optional[str] = None,
874
875
876
        init_method: Optional[Callable] = None,
        bias: bool = True,
        return_bias: bool = False,
877
        params_dtype: Optional[torch.dtype] = None,
878
        parallel_mode: Optional[str] = None,
cyanguwa's avatar
cyanguwa committed
879
        parameters_split: Optional[Union[Tuple[str, ...], Dict[str, int]]] = None,
880
        device: Union[torch.device, str] = "cuda",
881
        ub_overlap_ag: bool = False,
882
        ub_overlap_rs: bool = False,
883
        ub_overlap_rs_dgrad: bool = False,
884
885
        ub_bulk_dgrad: bool = False,
        ub_bulk_wgrad: bool = False,
886
        ub_name: Optional[str] = None,
887
        name: Optional[str] = None,
888
889
    ) -> None:
        super().__init__()
890
891

        params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype
892
893
894
895
896
897
        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
898
899
        self.get_rng_state_tracker = get_rng_state_tracker
        self.rng_tracker_name = rng_tracker_name
900
901
902
903
        self.name = name

        if TEDebugState.debug_enabled:
            self._turn_off_unsupported_features_in_debug()  # turn off userbuffers
904

905
906
        if device == "meta":
            assert parameters_split is None, "Cannot split module parameters on 'meta' device."
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
        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

928
        # Column parallel TP overlap options
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
        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
        )
947
948

        # Row parallel TP overlap options
949
950
951
952
953
954
        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
        )
955
956
957
958
959
960
961
962
963
964
965
966
967
968

        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

969
970
971
        # Initialize params in FP8
        with_fp8_params = FP8GlobalStateManager.with_fp8_parameters()

972
973
974
975
976
977
978
979
        # Contiguous buffers for params
        weight_tensor = torch.empty(
            self.out_features,
            self.in_features,
            device=device,
            dtype=params_dtype,
        )
        bias_tensor = None
980
        if self.use_bias:
981
982
983
984
985
            bias_tensor = torch.empty(
                self.out_features,
                device=device,
                dtype=params_dtype,
            )
986

987
988
989
990
        # Configure parameter splits
        self.weight_names = []
        self.bias_names = []
        self.parameter_split_sizes = []
991
        if parameters_split is None:
992
993
994
995
996
997
            # 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
998
        elif isinstance(parameters_split, dict):
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
            # 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
1011
        else:
1012
            raise TypeError("Invalid configuration for parameters split")
1013

1014
1015
1016
1017
1018
1019
        # 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}"
            )
1020

1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
        # 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

1031
1032
1033
1034
1035
        # 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.
1036
1037
1038
1039
1040
1041
1042
1043
        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)
1044
            if is_subview and with_fp8_params:
1045
1046
1047
                raise RuntimeError(
                    "Splitting QuantizedTensor into multiple params is not supported"
                )
1048

1049
            # Construct weight parameter
1050
1051
1052
1053
1054
1055
1056
            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,
            )
1057

1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
        # 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
1074

1075
        if with_fp8_params:
1076
1077
            self.init_fp8_metadata()

1078
        self.reset_parameters(defer_init=device == "meta")
1079

1080
1081
1082
1083
1084
1085
1086
        # 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

1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
    def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
        """Init scales and amaxes for fwd | bwd."""
        super().set_meta_tensor(fwd, recipe)

        # customize quantizers based on each recipe & layer configs
        recipe = FP8GlobalStateManager.get_fp8_recipe()
        if recipe.float8_current_scaling():
            self._customize_quantizers_float8_current_scaling(fwd, recipe)
        # elif for other recipes (mxfp8, etc.)

1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
    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)

1118
    @no_torch_dynamo()
1119
1120
1121
1122
    def forward(
        self,
        inp: torch.Tensor,
        is_first_microbatch: Optional[bool] = None,
1123
        fp8_output: Optional[bool] = False,
1124
        fp8_grad: Optional[bool] = False,
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
    ) -> 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)
        """
1147
1148
1149
1150
        debug = TEDebugState.debug_enabled
        if debug:
            self._validate_name()

1151
1152
1153
1154
        if FP8GlobalStateManager.fp8_graph_capturing():
            skip_fp8_weight_update = FP8GlobalStateManager.get_skip_fp8_weight_update_tensor()
        else:
            skip_fp8_weight_update = None
1155
1156
1157
        if skip_fp8_weight_update is not None:
            is_first_microbatch = False

1158
1159
1160
1161
1162
1163
1164
        if self.ub_overlap_rs_fprop:
            if get_ub(self.ub_name + "_fprop").is_fp8_ubuf():
                fp8_output = True
        if self.ub_overlap_rs_dgrad:
            if get_ub(self.ub_name + "_dgrad").is_fp8_ubuf():
                fp8_grad = True

1165
1166
        with self.prepare_forward(
            inp,
1167
            allow_non_contiguous=isinstance(inp, QuantizedTensor),
1168
        ) as inp:
1169
1170

            # Get concatenated weight and bias tensors
1171
            unfused_weights = [getattr(self, name) for name in self.weight_names]
1172
            if any(isinstance(w, QuantizedTensor) for w in unfused_weights):
1173
1174
1175
                if self.fp8:
                    if len(unfused_weights) != 1:
                        raise RuntimeError(
1176
                            "Splitting QuantizedTensor into multiple params is not supported"
1177
1178
                        )
                else:
1179
                    unfused_weights = [w.dequantize() for w in unfused_weights]
1180
            weight_tensor = noop_cat(unfused_weights)
1181
            if self.use_bias:
1182
                bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names])
1183
            else:
1184
1185
                bias_tensor = None

1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
            quantizers = (
                self._get_quantizers(fp8_output, fp8_grad)
                if not debug
                else self._get_debug_quantizers(fp8_output, fp8_grad)
            )
            if debug:
                if not any_feature_enabled(quantizers):
                    # If no feature is used, then run faster implementation with debug = False.
                    quantizers = self._get_quantizers(fp8_output, fp8_grad)
                    debug = False

                if isinstance(weight_tensor, QuantizedTensor):
                    raise RuntimeError("FP8 weights are not supported in debug mode.")

1200
1201
1202
1203
1204
            (
                input_quantizer,
                weight_quantizer,
                output_quantizer,
                grad_input_quantizer,
1205
1206
1207
                grad_weight_quantizer,
                grad_output_quantizer,
            ) = quantizers
1208
1209
1210
1211
1212
1213

            # Make sure weight tensor has correct quantizer
            # Note: Quantizer might have changed if quantization
            # recipe changed
            if weight_quantizer is not None and isinstance(weight_tensor, QuantizedTensor):
                weight_tensor._quantizer = weight_quantizer
1214

1215
1216
1217
1218
1219
1220
1221
1222
1223
            if torch.is_grad_enabled():
                linear_fn = _Linear.apply
                args = []
            else:
                linear_fn = _Linear.forward
                args = [None]
            args += (
                weight_tensor,
                inp,
1224
                bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None,
1225
1226
1227
                is_first_microbatch,
                self.fp8,
                self.fp8_calibration,
1228
1229
1230
1231
                input_quantizer,
                weight_quantizer,
                output_quantizer,
                grad_input_quantizer,
1232
1233
                grad_weight_quantizer,
                grad_output_quantizer,
1234
                self.fuse_wgrad_accumulation,
1235
                is_cpu_offload_enabled(),
1236
1237
1238
1239
1240
1241
1242
                self.tp_group,
                self.tp_size,
                self.sequence_parallel,
                self.tp_size > 1,
                self.activation_dtype,
                self.parallel_mode,
                torch.is_grad_enabled(),
1243
1244
1245
1246
1247
1248
                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,
1249
                self.ub_name,
1250
                fp8_output,
1251
                self.fsdp_group,
1252
1253
                self,
                skip_fp8_weight_update,
1254
                debug,
1255
1256
1257
1258
1259
1260
1261
1262
            )
            out = linear_fn(*args)
        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
1263
1264
1265

    def _get_quantizers(self, fp8_output, fp8_grad):
        if not self.fp8:
1266
            return [None] * 6
1267
        grad_input_quantizer = None
1268
        grad_weight_quantizer = None
1269
1270
1271
        grad_output_quantizer = None
        output_quantizer = None
        input_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT]
1272
        input_quantizer.internal = False
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
        weight_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT]
        weight_quantizer.internal = True
        if fp8_output:
            output_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT]
        if torch.is_grad_enabled():
            grad_output_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1]
            grad_output_quantizer.internal = True
            if fp8_grad:
                grad_input_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1]
        return (
            input_quantizer,
            weight_quantizer,
            output_quantizer,
            grad_input_quantizer,
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
            grad_weight_quantizer,
            grad_output_quantizer,
        )

    def _get_debug_quantizers(self, fp8_output, fp8_grad):
        original_quantizers = self._get_quantizers(fp8_output, fp8_grad)
        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)
1300
        )
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347

    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