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

5
import math
6
import os
7
from typing import Dict, List, Tuple, Optional
8
import pytest
9
import random
10
11
12
13

import torch
import torch.nn as nn
from torch.nn import Parameter
yuguo's avatar
yuguo committed
14
from torch.utils.cpp_extension import IS_HIP_EXTENSION
15

16
from transformer_engine.pytorch.quantization import FP8GlobalStateManager
17
18
19
from transformer_engine.pytorch.utils import (
    init_method_normal,
    scaled_init_method_normal,
20
21
22
    attention_mask_func,
)
from transformer_engine.pytorch import (
23
24
    autocast,
    quantized_model_init,
25
26
27
28
    DotProductAttention,
    LayerNormLinear,
    LayerNormMLP,
    Linear,
29
    GroupedLinear,
30
31
32
33
    MultiheadAttention,
    RMSNorm,
    TransformerLayer,
    LayerNorm,
34
35
    Fp8Padding,
    Fp8Unpadding,
36
37
    Float8Quantizer,
    Float8CurrentScalingQuantizer,
38
39
40
41
42
43
    MXFP8Quantizer,
    get_device_compute_capability,
    is_fp8_available,
    is_mxfp8_available,
    is_fp8_block_scaling_available,
    is_bf16_available,
44
    is_nvfp4_available,
45
)
46
from transformer_engine.pytorch import torch_version
47
from transformer_engine.pytorch import checkpoint as te_checkpoint
48
49
from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm
from transformer_engine.common import recipe
50
import transformer_engine_torch as tex
51
from utils import ModelConfig, reset_rng_states
52

53

54
# Only run FP8 tests on supported devices.
55
56
fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True)
mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True)
wenjh's avatar
wenjh committed
57
fp8_block_scaling_available, reason_for_no_fp8_block_scaling = is_fp8_block_scaling_available(return_reason=True)
58
nvfp4_available = is_nvfp4_available()
59

60
sm_80plus = get_device_compute_capability() >= (8, 0)
61

62
seed = 1234
63
64
# Reset RNG states.
reset_rng_states()
65

66
67
68
69
if torch_version() >= (2, 7, 0):
    torch._dynamo.config.recompile_limit = 16
else:
    torch._dynamo.config.cache_size_limit = 16
70
71
72


model_configs = {
73
74
    "small": ModelConfig(1, 128, 8, 16, num_layers=4),
    "126m": ModelConfig(1, 2048, 12, 64, num_layers=12),
75
}
76
model_configs_inference = {
77
    "126m": ModelConfig(1, 256, 12, 64, num_layers=12),
78
}
79
backends_inference = ["FlashAttention", "UnfusedAttention", "FusedAttention"]
80
81
82
module_inference = ["TransformerLayer", "MultiheadAttention"]
input_formats_inference = ["sbhd", "bshd"]

83
param_types = [torch.float32, torch.float16]
84
if is_bf16_available():  # bf16 requires sm_80 or higher
85
86
87
88
89
90
    param_types.append(torch.bfloat16)

batch_sizes = [1, 2]

all_boolean = [True, False]

91
92
93
94
95
96
97
98
99
100
101
102
all_activations = [
    "gelu",
    "geglu",
    "qgelu",
    "qgeglu",
    "relu",
    "reglu",
    "srelu",
    "sreglu",
    "silu",
    "swiglu",
]
103

104
105
all_normalizations = ["LayerNorm", "RMSNorm"]

106
107
mask_types = ["causal", "no_mask"]

108
NVTE_TEST_NVINSPECT_ENABLED = int(os.environ.get("NVTE_TEST_NVINSPECT_ENABLED", "0"))
109
110
111
112
113
114
115
116
117
118
119
120
121

if NVTE_TEST_NVINSPECT_ENABLED:
    # The numerics of all the layers should work the same,
    # when debug=True. I fed them with dummy feature
    # to prevent switching off debug, which can happen if
    # no feature is active.
    import nvdlfw_inspect.api as debug_api

    debug_api.initialize(
        os.environ["NVTE_TEST_NVINSPECT_CONFIG_FILE"],
        feature_dirs=os.environ["NVTE_TEST_NVINSPECT_FEATURE_DIRS"],
    )

122

123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def nvfp4_rht_and_2d_quantization():
    nvfp4_recipe = recipe.NVFP4BlockScaling()
    nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams(
        random_hadamard_transform=True, fp4_2d_quantization=False
    )
    nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams(
        random_hadamard_transform=False, fp4_2d_quantization=True
    )
    nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams(
        random_hadamard_transform=True, fp4_2d_quantization=False
    )
    return nvfp4_recipe


def check_rht_usage(recipe: recipe.Recipe) -> bool:
    # if using RHT, we can only support bf16
    # check fp4_quant_fwd_inp, fp4_quant_fwd_weight, fp4_quant_bwd_grad
    if recipe.nvfp4():
        if (
            recipe.fp4_quant_fwd_inp.random_hadamard_transform
            or recipe.fp4_quant_fwd_weight.random_hadamard_transform
            or recipe.fp4_quant_bwd_grad.random_hadamard_transform
        ):
            return True
    return False


def get_nvfp4_inp_supported_dtypes(recipe: recipe.Recipe, dtype: torch.dtype) -> bool:
    supported_input_dtypes = []
    if recipe.nvfp4():
        supported_input_dtypes.append(torch.bfloat16)
        # if not using RHT, we can add fp32 as well
    if not check_rht_usage(recipe):
        supported_input_dtypes.append(torch.float32)
    return supported_input_dtypes


160
161
162
163
164
165
166
167
fp8_recipes = []
if mxfp8_available:
    fp8_recipes.append(recipe.MXFP8BlockScaling())
if fp8_block_scaling_available:
    fp8_recipes.append(recipe.Float8BlockScaling())
if fp8_available:
    fp8_recipes.append(recipe.Float8CurrentScaling())
    fp8_recipes.append(recipe.DelayedScaling())
168
169
if nvfp4_available:
    fp8_recipes.append(nvfp4_rht_and_2d_quantization())
170

171
172
173
174
175
use_cutlass_grouped_gemm = [False]
# Only enable cutlass grouped gemm on Hopper
if torch.cuda.get_device_capability() == (9, 0):
    use_cutlass_grouped_gemm.append(True)

176

177
178
179
180
def get_causal_attn_mask(sq: int) -> torch.Tensor:
    return torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool()


181
182
def dtype_tols(dtype: torch.dtype) -> Dict[str, float]:
    """Estimated numerical error for a datatype
183

184
    Based on tolerances for torch.testing.assert_close.
185

186
187
188
189
190
191
192
193
194
195
196
    """
    if dtype == torch.float32:
        return dict(rtol=1.3e-6, atol=1e-5)
    if dtype == torch.float16:
        return dict(rtol=1e-3, atol=1e-5)
    if dtype == torch.bfloat16:
        return dict(rtol=1.6e-2, atol=1e-5)
    raise ValueError(f"Unsuppored dtype ({dtype})")


def assert_allclose(
197
    l1: List[torch.Tensor], l2: List[torch.Tensor], atol: float = None, rtol: float = None
198
) -> bool:
199
200
    """Ensures two lists are equal."""
    assert len(l1) == len(l2), "Unequal number of outputs."
201
    for i, (t1, t2) in enumerate(zip(l1, l2)):
202
        tols = dtype_tols(t2.dtype)
203
204
        if rtol is not None:
            tols["rtol"] = rtol
205
206
        if atol is not None:
            tols["atol"] = atol
207
        result = torch.allclose(t1, t2, **tols)
208
        if not result:
209
            diff = torch.abs(t1 - t2)
210
            tol = tols["atol"] + (tols["rtol"] * torch.abs(t2))
211
212
213
214
215
216
217
218
219
220
221
222
            exceed_mask = diff > tol
            if exceed_mask.any():
                indices = torch.nonzero(exceed_mask, as_tuple=True)
                max_diff = diff[exceed_mask].max()
                max_idx = (diff[exceed_mask] == max_diff).nonzero(as_tuple=True)[0][0]
                max_location = [idx[max_idx].item() for idx in indices]
                msg = (
                    f"Outputs not close enough in tensor at idx={i}. "
                    f"Maximum difference at location {max_location} "
                    f"with {t1[exceed_mask][max_idx].item()} vs {t2[exceed_mask][max_idx].item()} "
                    f"(diff {max_diff.item()})."
                )
223
            raise AssertionError(msg)
224
225


226
227
228
229
@pytest.fixture(autouse=True)
def reset_global_fp8_state():
    yield
    FP8GlobalStateManager.reset()
230
231


232
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
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
class TorchScaledMaskedSoftmax(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(
        self, inp: torch.Tensor, mask: torch.Tensor, scale: Optional[float] = None
    ) -> torch.Tensor:
        dtype = inp.dtype
        inp = inp.float()

        if scale is not None:
            inp = inp * scale
        mask_output = attention_mask_func(inp, mask) if mask is not None else inp

        probs = torch.nn.Softmax(dim=-1)(mask_output)
        probs = probs.to(dtype)
        return probs


class TorchDotProductAttention(torch.nn.Module):
    def __init__(
        self,
        kv_channels: int,
        attention_dropout: float = 0.0,
    ) -> None:
        super().__init__()

        self.norm_factor = math.sqrt(kv_channels)
        self.scale_mask_softmax = TorchScaledMaskedSoftmax()
        self.attention_dropout = torch.nn.Dropout(attention_dropout)

    def forward(
        self,
        query_layer: torch.Tensor,
        key_layer: torch.Tensor,
        value_layer: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        batch_size, seqlen = query_layer.shape[1], query_layer.shape[0]

        # [b, np, sq, sk]
        output_size = (
            query_layer.size(1),
            query_layer.size(2),
            query_layer.size(0),
            key_layer.size(0),
        )

        # [sq, b, np, hn] -> [sq, b * np, hn]
281
        query_layer = query_layer.reshape(output_size[2], output_size[0] * output_size[1], -1)
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
        # [sk, b, np, hn] -> [sk, b * np, hn]
        key_layer = key_layer.reshape(output_size[3], output_size[0] * output_size[1], -1)

        # preallocting result tensor: [b * np, sq, sk]
        matmul_result = torch.empty(
            output_size[0] * output_size[1],
            output_size[2],
            output_size[3],
            dtype=query_layer.dtype,
            device=torch.cuda.current_device(),
        )

        # Raw attention scores. [b * np, sq, sk]
        matmul_result = torch.baddbmm(
            matmul_result,
            query_layer.transpose(0, 1),  # [b * np, sq, hn]
            key_layer.transpose(0, 1).transpose(1, 2),  # [b * np, hn, sk]
            beta=0.0,
            alpha=(1.0 / self.norm_factor),
        )

        # change view to [b, np, sq, sk]
        attention_scores = matmul_result.view(*output_size)

        # attention scores and attention mask [b, np, sq, sk]
        attention_probs = self.scale_mask_softmax(attention_scores, attention_mask)
        attention_probs = self.attention_dropout(attention_probs)

        # value_layer -> context layer.
        # [sk, b, np, hn] --> [b, np, sq, hn]
        output_size = (
            value_layer.size(1),
            value_layer.size(2),
            query_layer.size(0),
            value_layer.size(3),
        )

        # change view [sk, b * np, hn]
320
        value_layer = value_layer.reshape(value_layer.size(0), output_size[0] * output_size[1], -1)
321
322

        # change view [b * np, sq, sk]
323
        attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338

        # matmul: [b * np, sq, hn]
        context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))

        # change view [b, np, sq, hn]
        context_layer = context_layer.view(*output_size)

        # [b, np, sq, hn] --> [sq, b, np, hn]
        context_layer = context_layer.permute(2, 0, 1, 3).contiguous()

        # [sq, b, np, hn] --> [sq, b, hp]
        context_layer = context_layer.view(seqlen, batch_size, -1)

        return context_layer

339

340
class TorchLayerNorm(nn.Module):
341
    def __init__(self, in_features: int, eps: float, zero_centered_gamma: bool):
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
        super().__init__()
        self.eps = eps
        self.in_features = in_features
        self.zero_centered_gamma = zero_centered_gamma

        initial_value = torch.ones(in_features) if zero_centered_gamma else torch.zeros(in_features)
        self.weight = nn.Parameter(initial_value)
        self.bias = nn.Parameter(torch.zeros(in_features))
        self.register_parameter("weight", self.weight)
        self.register_parameter("bias", self.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.weight if not self.zero_centered_gamma else 1 + self.weight
        w = w.to(torch.float32)
        b = self.bias.to(torch.float32)
        inp = x.to(torch.float32)
358
359
360
        out = torch.nn.functional.layer_norm(
            inp, (self.in_features,), weight=w, bias=b, eps=self.eps
        )
361
362
        return out.to(x.dtype)

363

364
365
# Adapted from https://github.com/bzhangGo/rmsnorm/blob/c6691f20ec0af4128c8159c903071f7575404295/rmsnorm_torch.py
class TorchRMSNorm(nn.Module):
366
    def __init__(self, in_features, zero_centered_gamma, eps=1e-5):
367
368
369
370
        super().__init__()

        self.eps = eps
        self.in_features = in_features
371
        self.zero_centered_gamma = zero_centered_gamma
372

373
374
        initial_value = torch.ones(in_features) if zero_centered_gamma else torch.zeros(in_features)
        self.weight = nn.Parameter(initial_value)
375
376
377
        self.register_parameter("weight", self.weight)

    def forward(self, x):
378
        norm_x2 = torch.sum(x.float() ** 2, dim=-1, keepdim=True)
379
380
        d_x = self.in_features

381
        rms_x2 = norm_x2 / d_x + self.eps
382
        r_rms_x = rms_x2 ** (-1.0 / 2)
383
        x_normed = x * r_rms_x
384

385
386
387
388
        w = self.weight.float()
        if self.zero_centered_gamma:
            w = 1 + w
        return (w * x_normed).to(x.dtype)
389

390

391
class TorchLayerNormLinear(nn.Module):
392
393
394
395
396
397
398
    def __init__(
        self,
        in_features: int,
        out_features: int,
        eps: float,
        normalization: str = "LayerNorm",
        zero_centered_gamma: bool = False,
399
        bias: bool = True,
400
    ):
401
        super().__init__()
402
        if normalization == "LayerNorm":
403
404
405
            self.layernorm = TorchLayerNorm(
                in_features, eps=eps, zero_centered_gamma=zero_centered_gamma
            )
406
        elif normalization == "RMSNorm":
407
408
409
            self.layernorm = TorchRMSNorm(
                in_features, eps=eps, zero_centered_gamma=zero_centered_gamma
            )
410
411
412
        else:
            raise RuntimeError("Unsupported normalization")

413
        self.linear = nn.Linear(in_features, out_features, bias=bias)
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear(self.layernorm(x))


class TorchMHA(nn.Module):
    def __init__(self, hidden_size: int, num_attention_heads: int):
        super().__init__()
        self.mhsa = nn.MultiheadAttention(
            embed_dim=hidden_size,
            num_heads=num_attention_heads,
            dropout=0.1,
            bias=True,
            batch_first=False,
        )

430
431
    def forward(self, x, attention_mask=None):
        output = self.mhsa(x, x, x, attn_mask=attention_mask, need_weights=False)
432
433
434
435
        if isinstance(output, tuple):
            output = output[0]
        return output

436

437
438
439
class TorchQuickGELU(nn.Module):
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return input * torch.sigmoid(1.702 * input)
440

441

442
443
444
445
class TorchSquaredRELU(nn.Module):
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return (input > 0) * input * input

446

447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
class TorchGroupedLinearWithPadding(nn.Module):

    def __init__(
        self, num_gemms, in_features, out_features, bias, params_dtype, parallel_mode, fp8
    ) -> None:
        super().__init__()

        self.padding = Fp8Padding(num_gemms)
        self.linear_fn = GroupedLinear(
            num_gemms,
            in_features,
            out_features,
            bias=bias,
            params_dtype=params_dtype,
            parallel_mode=parallel_mode,
            device="cuda",
        )
        self.unpadding = Fp8Unpadding(num_gemms)

        self.fp8 = fp8

    def forward(self, inp: torch.Tensor, m_splits: List[int]) -> torch.Tensor:
        if self.fp8:
            orig_m_splits = m_splits
            inp, m_splits = self.padding(inp, m_splits)

        out = self.linear_fn(inp, m_splits)

        if self.fp8:
            out = self.unpadding(out, orig_m_splits)

        return out


481
482
_supported_act = {
    "gelu": nn.GELU(approximate="tanh"),
483
    "geglu": nn.GELU(approximate="tanh"),
484
    "qgelu": TorchQuickGELU(),
485
486
487
    "qgeglu": TorchQuickGELU(),
    "relu": nn.ReLU(),
    "reglu": nn.ReLU(),
488
    "srelu": TorchSquaredRELU(),
489
490
491
    "sreglu": TorchSquaredRELU(),
    "silu": nn.SiLU(),
    "swiglu": nn.SiLU(),
492
}
493

494

495
496
497
498
499
500
501
class TorchGLU(nn.Module):
    def __init__(self, activation: str):
        super().__init__()
        self.act = _supported_act[activation]

    def forward(self, x):
        shape = x.size(-1)
502
503
        a = x[..., : shape // 2]
        b = x[..., (shape // 2) :]
504
505
        a = self.act(a)
        return a * b
506

507

508
class TorchLayerNormMLP(nn.Module):
509
510
511
512
513
514
515
    def __init__(
        self,
        hidden_size: int,
        ffn_hidden_size: int,
        eps: float = 1e-5,
        activation="gelu",
        normalization: str = "LayerNorm",
516
        bias: bool = True,
517
    ):
518
        super().__init__()
519
        if normalization == "LayerNorm":
520
            self.ln = TorchLayerNorm(hidden_size, eps=eps, zero_centered_gamma=False)
521
        elif normalization == "RMSNorm":
522
            self.ln = TorchRMSNorm(hidden_size, eps=eps, zero_centered_gamma=False)
523
524
        else:
            raise RuntimeError("Unsupported normalization")
525
        if "glu" in activation:
526
527
528
529
530
531
            fc1_output_features = 2 * ffn_hidden_size
            self.gelu = TorchGLU(activation)
        else:
            fc1_output_features = ffn_hidden_size
            self.gelu = _supported_act[activation]

532
533
        self.fc1 = nn.Linear(hidden_size, fc1_output_features, bias=bias)
        self.fc2 = nn.Linear(ffn_hidden_size, hidden_size, bias=bias)
534
535

    def forward(self, x):
536
537
        t = self.gelu(self.fc1(self.ln(x)))
        return self.fc2(t)
538
539
540


class TorchGPT(nn.Module):
541
542
543
    def __init__(
        self, hidden_size: int, eps: float, num_attention_heads: int, parallel_attention_mlp: bool
    ):
544
        super().__init__()
545
        self.ln = nn.LayerNorm(hidden_size, eps=eps)
546
        self.causal_attn = TorchMHA(hidden_size, num_attention_heads)
547
        self.ln_mlp = TorchLayerNormMLP(hidden_size, 4 * hidden_size, eps)
548
        self.parallel_attention_mlp = parallel_attention_mlp
549
550
551
552

    def forward(
        self,
        x: torch.Tensor,
553
        attention_mask: Optional[torch.Tensor] = None,
554
    ) -> torch.Tensor:
555
        a = self.ln(x)
556
        b = self.causal_attn(a, attention_mask)
557
558
559
560
561
562
563
        if self.parallel_attention_mlp:
            n = self.ln_mlp(x)
            x = x + nn.functional.dropout(b + n, p=0.1, training=self.training)
        else:
            x = x + nn.functional.dropout(b, p=0.1, training=self.training)
            n = self.ln_mlp(x)
            x = x + nn.functional.dropout(n, p=0.1, training=self.training)
564
565
566
        return x


567
568
569
def _test_e2e_selective_recompute(
    bs, dtype, config, fp8, recipe, fp8_model_params=False, recompute=False
):
570
    reset_rng_states()
571
    FP8GlobalStateManager.reset()
572
573
574
575
576

    sigma = 0.023
    init_method = init_method_normal(sigma)
    output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)

577
    with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe):
578
579
580
        block = TransformerLayer(
            config.hidden_size,
            4 * config.hidden_size,
581
            config.num_heads,
582
583
584
585
586
            layernorm_epsilon=config.eps,
            init_method=init_method,
            output_layer_init_method=output_layer_init_method,
            hidden_dropout=0.1,
            attention_dropout=0.1,
587
            kv_channels=config.kv_channels,
588
589
590
591
592
            apply_residual_connection_post_layernorm=False,
            output_layernorm=False,
            params_dtype=dtype,
            fuse_qkv_params=True,
            device="cuda",
593
594
595
        )

    te_inp_hidden_states = torch.randn(
596
        (config.max_seqlen_q, bs, config.hidden_size),
597
598
599
600
        dtype=dtype,
        device="cuda",
        requires_grad=True,
    )
601
    te_inp_hidden_states.retain_grad()
602
    te_inp_attn_mask = get_causal_attn_mask(config.max_seqlen_q)
603

604
    with autocast(enabled=fp8, recipe=recipe):
605
606
        te_out = block(
            te_inp_hidden_states,
607
            attention_mask=te_inp_attn_mask,
608
            checkpoint_core_attention=recompute,
609
610
611
612
613
614
615
616
617
618
619
620
621
622
        )
    loss = te_out.sum()
    loss.backward()
    torch.cuda.synchronize()

    outputs = [te_out, te_inp_hidden_states.grad]
    for p in block.parameters():
        if p.requires_grad:
            outputs.append(p.grad)
    return outputs


@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
623
@pytest.mark.parametrize("model", ["126m"])
624
@pytest.mark.parametrize("fp8", all_boolean)
625
@pytest.mark.parametrize("recipe", fp8_recipes)
626
@pytest.mark.parametrize("fp8_model_params", all_boolean)
627
def test_gpt_selective_activation_recompute(dtype, bs, model, fp8, recipe, fp8_model_params):
628
629
    if fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED:
        pytest.skip("FP8 parameters are not supported in debug mode.")
630
631
632
    if fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

633
634
635
636
637
    if fp8 and recipe.nvfp4():
        if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype):
            pytest.skip(
                f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}"
            )
638

639
640
    config = model_configs[model]

641
    outputs = _test_e2e_selective_recompute(
642
        bs, dtype, config, fp8, recipe, fp8_model_params, recompute=False
643
644
    )
    outputs_recompute = _test_e2e_selective_recompute(
645
        bs, dtype, config, fp8, recipe, fp8_model_params, recompute=True
646
    )
647
648
649
650
651
652
653

    # Check that results match
    tols = dtype_tols(dtype)
    if dtype in (torch.float16, torch.bfloat16):
        tols["atol"] = 1e-4
    if fp8 or fp8_model_params:
        tols.update(dict(rtol=0.125, atol=0.0675))
654

655
656
657
658
659
660
661
    for i, (ref, test) in enumerate(zip(outputs, outputs_recompute)):
        torch.testing.assert_close(
            test,
            ref,
            msg=f"Mismatch in tensor {i}",
            **tols,
        )
662
663


664
def _test_e2e_full_recompute(
665
    bs, dtype, config, fp8, recipe, fp8_model_params=False, recompute=False, use_reentrant=True
666
):
667
668
669
    reset_rng_states()
    FP8GlobalStateManager.reset()

670
671
672
673
    sigma = 0.023
    init_method = init_method_normal(sigma)
    output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)

674
    with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe):
675
        block = TransformerLayer(
676
677
            config.hidden_size,
            4 * config.hidden_size,
678
            config.num_heads,
679
680
681
682
683
            layernorm_epsilon=config.eps,
            init_method=init_method,
            output_layer_init_method=output_layer_init_method,
            hidden_dropout=0.1,
            attention_dropout=0.1,
684
            kv_channels=config.kv_channels,
685
686
687
            apply_residual_connection_post_layernorm=False,
            output_layernorm=False,
            params_dtype=dtype,
688
            fuse_qkv_params=True,
689
            device="cuda",
690
        )
691

692
    te_inp_hidden_states = torch.randn(
693
        (config.max_seqlen_q, bs, config.hidden_size),
694
695
696
697
        dtype=dtype,
        device="cuda",
        requires_grad=use_reentrant,
    )
698
699
    if use_reentrant:
        te_inp_hidden_states.retain_grad()
700
    te_inp_attn_mask = get_causal_attn_mask(config.max_seqlen_q)
701

702
    with autocast(enabled=fp8, recipe=recipe):
703
704
705
706
707
708
        if recompute:
            te_out = te_checkpoint(
                block,
                te_inp_hidden_states,
                attention_mask=te_inp_attn_mask,
                checkpoint_core_attention=False,
709
710
711
                distribute_saved_activations=False,
                tp_group=None,
                use_reentrant=use_reentrant,
712
713
714
715
716
717
718
719
720
721
722
            )
        else:
            te_out = block(
                te_inp_hidden_states,
                attention_mask=te_inp_attn_mask,
                checkpoint_core_attention=False,
            )
    loss = te_out.sum()
    loss.backward()
    torch.cuda.synchronize()

723
724
725
726
727
728
    outputs = [te_out]
    names = ["output"]
    if use_reentrant:
        outputs.append(te_inp_hidden_states.grad)
        names.append("input")
    for name, p in block.named_parameters():
729
730
        if p.requires_grad:
            outputs.append(p.grad)
731
732
733
            names.append(name)

    return outputs, names
734
735
736
737


@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
738
@pytest.mark.parametrize("model", ["126m"])
739
@pytest.mark.parametrize("fp8", all_boolean)
740
@pytest.mark.parametrize("recipe", fp8_recipes)
741
@pytest.mark.parametrize("fp8_model_params", all_boolean)
742
@pytest.mark.parametrize("use_reentrant", all_boolean)
743
744
745
def test_gpt_full_activation_recompute(
    dtype, bs, model, fp8, recipe, fp8_model_params, use_reentrant
):
746
747
    if fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)
748
749
    if fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED:
        pytest.skip("FP8 parameters are not supported in debug mode.")
750

751
752
753
754
755
    if fp8 and recipe.nvfp4():
        if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype):
            pytest.skip(
                f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}"
            )
756
757
758

    config = model_configs[model]

759
760
761
762
    if not use_reentrant:
        # Non-reentrant checkpoint becomes non-deterministic with bias+GELU fusion
        os.environ["NVTE_BIAS_GELU_NVFUSION"] = "0"

763
    outputs, names = _test_e2e_full_recompute(
764
765
766
767
768
769
770
771
        bs,
        dtype,
        config,
        fp8,
        recipe,
        fp8_model_params,
        recompute=False,
        use_reentrant=use_reentrant,
772
773
    )
    outputs_recompute, _ = _test_e2e_full_recompute(
774
775
776
777
778
779
780
781
        bs,
        dtype,
        config,
        fp8,
        recipe,
        fp8_model_params,
        recompute=True,
        use_reentrant=use_reentrant,
782
    )
783
784
785
786
787

    if not use_reentrant:
        # Reset bias+GELU fusion flag to avoid contaminating other tests
        del os.environ["NVTE_BIAS_GELU_NVFUSION"]

788
789
790
791
792
793
794
795
796
797
798
799
800
    # Check that results match
    tols = dtype_tols(dtype)
    if dtype in (torch.float16, torch.bfloat16):
        tols["atol"] = 1e-3
    if fp8 or fp8_model_params:
        tols.update(dict(rtol=0.125, atol=0.0675))
    for i, (ref, test) in enumerate(zip(outputs, outputs_recompute)):
        torch.testing.assert_close(
            test,
            ref,
            msg=f"Mismatch in tensor {i}",
            **tols,
        )
801
802
803
804
805
806


def _test_e2e_checkpointing_get_model(config, dtype):
    sigma = 0.023
    init_method = init_method_normal(sigma)
    output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
807

808
809
810
    return TransformerLayer(
        config.hidden_size,
        4 * config.hidden_size,
811
        config.num_heads,
812
813
814
815
816
        layernorm_epsilon=config.eps,
        init_method=init_method,
        output_layer_init_method=output_layer_init_method,
        hidden_dropout=0.1,
        attention_dropout=0.1,
817
        kv_channels=config.kv_channels,
818
819
820
821
        apply_residual_connection_post_layernorm=False,
        output_layernorm=False,
        params_dtype=dtype,
        device="cuda",
822
823
824
825
826
827
828
    )


def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path="checkpoint.pt"):
    reset_rng_states()

    te_inp_hidden_states = torch.randn(
829
        (config.max_seqlen_q, bs, config.hidden_size),
830
831
832
833
        dtype=dtype,
        device="cuda",
        requires_grad=True,
    )
834
835
836
837
838
839
840
    te_inp_hidden_states.retain_grad()

    block = _test_e2e_checkpointing_get_model(config, dtype)

    for _ in range(steps // 2):
        te_out = block(
            te_inp_hidden_states,
841
            None,
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
        )
        loss = te_out.sum()
        loss.backward()

    if checkpoint:
        # This process is necessary so that we can start afresh with
        # a new model while erasing all internal state to ensure that
        # loading from a checkpoint gives bitwise identical results.
        # Since gradients are being accumulated, it is important to
        # restore them post loading the checkpoint.
        torch.save(block.state_dict(), path)

        param_grads = []
        for p in block.parameters():
            if p.requires_grad:
                param_grads.append(p.grad.clone())

859
860
861
        _cpu_rng_state = torch.get_rng_state()
        _cuda_rng_state = torch.cuda.get_rng_state()

862
863
        del block
        block = _test_e2e_checkpointing_get_model(config, dtype)
864
        block.load_state_dict(torch.load(path, weights_only=False))
865
866
        torch.set_rng_state(_cpu_rng_state)
        torch.cuda.set_rng_state(_cuda_rng_state)
867
868
869
870
871
872
873
874
875
876

        for p in block.parameters():
            if p.requires_grad:
                p.grad = param_grads.pop(0)

        assert not param_grads, "Oops!"

    for _ in range(steps // 2):
        te_out = block(
            te_inp_hidden_states,
877
            None,
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
        )
        loss = te_out.sum()
        loss.backward()

    torch.cuda.synchronize()

    if os.path.exists(path):
        os.remove(path)

    outputs = [te_out, te_inp_hidden_states.grad]
    for p in block.parameters():
        if p.requires_grad:
            outputs.append(p.grad)
    return outputs


@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
896
@pytest.mark.parametrize("model", ["126m"])
897
898
899
def test_gpt_checkpointing(dtype, bs, model):
    config = model_configs[model]
    outputs = _test_e2e_checkpointing(bs, dtype, config, checkpoint=False)
900
    outputs_checkpoint = _test_e2e_checkpointing(bs, dtype, config, checkpoint=True)
901
902
903
904
905
906
907
908
909
910
911
912

    # Check that results match
    tols = dtype_tols(dtype)
    if dtype in (torch.float16, torch.bfloat16):
        tols.update(dict(rtol=2e-2, atol=2e-3))
    for i, (ref, test) in enumerate(zip(outputs, outputs_checkpoint)):
        torch.testing.assert_close(
            test,
            ref,
            msg=f"Mismatch in tensor {i}",
            **tols,
        )
913
914
915
916
917
918


def _test_e2e_gpt_accuracy(block, bs, dtype, config):
    reset_rng_states()

    inp_hidden_states = torch.randn(
919
        (config.max_seqlen_q, bs, config.hidden_size),
920
921
922
923
        dtype=dtype,
        device="cuda",
        requires_grad=True,
    )
924
    inp_hidden_states.retain_grad()
925
    inp_attn_mask = get_causal_attn_mask(config.max_seqlen_q)
926

927
    out = block(inp_hidden_states, attention_mask=inp_attn_mask)
928
929
930
931
932
933
934
935
936
937
938
939
940
    loss = out.sum()
    loss.backward()

    torch.cuda.synchronize()
    outputs = [out, inp_hidden_states.grad]
    for p in block.parameters():
        if p.requires_grad:
            outputs.append(p.grad)
    return outputs


@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
941
@pytest.mark.parametrize("model", ["small"])
942
943
@pytest.mark.parametrize("parallel_attention_mlp", all_boolean)
def test_gpt_accuracy(dtype, bs, model, parallel_attention_mlp):
944
945
    config = model_configs[model]

946
947
948
    te_gpt = TransformerLayer(
        hidden_size=config.hidden_size,
        ffn_hidden_size=4 * config.hidden_size,
949
        num_attention_heads=config.num_heads,
950
951
952
953
954
955
956
957
958
        layernorm_epsilon=config.eps,
        attention_dropout=0.1,
        hidden_dropout=0.1,
        params_dtype=dtype,
        fuse_qkv_params=True,
        qkv_weight_interleaved=False,
        parallel_attention_mlp=parallel_attention_mlp,
        device="cuda",
    ).eval()
959
960
961
962
963

    torch_gpt = (
        TorchGPT(
            config.hidden_size,
            config.eps,
964
            config.num_heads,
965
            parallel_attention_mlp=parallel_attention_mlp,
966
967
968
969
970
971
972
973
        )
        .to(dtype=dtype)
        .cuda()
        .eval()
    )

    # Share params
    with torch.no_grad():
974
        torch_gpt.ln.weight = Parameter(
975
976
            te_gpt.self_attention.layernorm_qkv.layer_norm_weight.clone()
        )
977
        torch_gpt.ln.bias = Parameter(te_gpt.self_attention.layernorm_qkv.layer_norm_bias.clone())
978
979
980
981
982
983
984
985
986
987
988
989
        torch_gpt.causal_attn.mhsa.in_proj_weight = Parameter(
            te_gpt.self_attention.layernorm_qkv.weight.clone()
        )
        torch_gpt.causal_attn.mhsa.in_proj_bias = Parameter(
            te_gpt.self_attention.layernorm_qkv.bias.clone()
        )
        torch_gpt.causal_attn.mhsa.out_proj.weight = Parameter(
            te_gpt.self_attention.proj.weight.clone()
        )
        torch_gpt.causal_attn.mhsa.out_proj.bias = Parameter(
            te_gpt.self_attention.proj.bias.clone()
        )
990
991
992
993
994
995
        torch_gpt.ln_mlp.ln.weight = Parameter(te_gpt.layernorm_mlp.layer_norm_weight.clone())
        torch_gpt.ln_mlp.ln.bias = Parameter(te_gpt.layernorm_mlp.layer_norm_bias.clone())
        torch_gpt.ln_mlp.fc1.weight = Parameter(te_gpt.layernorm_mlp.fc1_weight.clone())
        torch_gpt.ln_mlp.fc1.bias = Parameter(te_gpt.layernorm_mlp.fc1_bias.clone())
        torch_gpt.ln_mlp.fc2.weight = Parameter(te_gpt.layernorm_mlp.fc2_weight.clone())
        torch_gpt.ln_mlp.fc2.bias = Parameter(te_gpt.layernorm_mlp.fc2_bias.clone())
996
997
998
999

    te_outputs = _test_e2e_gpt_accuracy(te_gpt, bs, dtype, config)
    torch_outputs = _test_e2e_gpt_accuracy(torch_gpt, bs, dtype, config)

1000
1001
1002
1003
1004
1005
    atol = {
        torch.float32: 5e-3,
        torch.half: 5e-2,
        torch.bfloat16: 1e-1,
    }

1006
    # Check output.
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
    assert_allclose(te_outputs[0], torch_outputs[0], atol[dtype])

    # Check gradients, only for small model
    if model == "small":
        atol[torch.float32] = 5e-2
        rtol = {
            torch.float32: 1e-2,
            torch.half: 1e-2,
            torch.bfloat16: 1e-2,
        }
        for te_output, torch_output in zip(te_outputs[1:], torch_outputs[1:]):
            assert_allclose(te_output, torch_output, atol[dtype], rtol[dtype])
1019
1020


1021
def _test_mha_accuracy(block, bs, dtype, config, mask_type, te=True):
1022
1023
1024
    reset_rng_states()

    inp_hidden_states = torch.randn(
1025
        (config.max_seqlen_q, bs, config.hidden_size),
1026
1027
1028
1029
        dtype=dtype,
        device="cuda",
        requires_grad=True,
    )
1030
    inp_hidden_states.retain_grad()
1031
    inp_attn_mask = get_causal_attn_mask(config.max_seqlen_q) if mask_type == "causal" else None
1032

1033
1034
1035
1036
1037
1038
    forward_kwargs = {}
    if te:
        forward_kwargs["attn_mask_type"] = mask_type
    forward_kwargs["attention_mask"] = inp_attn_mask

    out = block(inp_hidden_states, **forward_kwargs)
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
    loss = out.sum()
    loss.backward()

    torch.cuda.synchronize()
    outputs = [out, inp_hidden_states.grad]
    for p in block.parameters():
        if p.requires_grad:
            outputs.append(p.grad)
    return outputs


@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
1052
@pytest.mark.parametrize("model", ["small"])
1053
1054
1055
1056
@pytest.mark.parametrize("mask_type", mask_types)
def test_mha_accuracy(dtype, bs, model, mask_type):
    config = model_configs[model]

1057
1058
    te_mha = MultiheadAttention(
        config.hidden_size,
1059
        config.num_heads,
1060
1061
1062
1063
1064
1065
        fuse_qkv_params=True,
        params_dtype=dtype,
        qkv_weight_interleaved=False,
        input_layernorm=False,
        device="cuda",
    ).eval()
1066
1067
1068
1069

    torch_mha = (
        TorchMHA(
            config.hidden_size,
1070
            config.num_heads,
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
        )
        .to(dtype=dtype)
        .cuda()
        .eval()
    )

    # Share params
    with torch.no_grad():
        torch_mha.mhsa.in_proj_weight = Parameter(te_mha.qkv.weight.clone())
        torch_mha.mhsa.in_proj_bias = Parameter(te_mha.qkv.bias.clone())
        torch_mha.mhsa.out_proj.weight = Parameter(te_mha.proj.weight.clone())
        torch_mha.mhsa.out_proj.bias = Parameter(te_mha.proj.bias.clone())

1084
1085
    te_outputs = _test_mha_accuracy(te_mha, bs, dtype, config, mask_type, te=True)
    torch_outputs = _test_mha_accuracy(torch_mha, bs, dtype, config, mask_type, te=False)
1086
1087
1088
1089
1090
1091
1092

    # Check output.
    if dtype == torch.float32:
        assert_allclose(te_outputs[0], torch_outputs[0], 5e-3)
    else:
        assert_allclose(te_outputs[0], torch_outputs[0], 5e-2)

1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
    # Check gradients, only for small model
    if model == "small":
        atol = {
            torch.float32: 5e-2,
            torch.half: 5e-2,
            torch.bfloat16: 5e-2,
        }
        rtol = {
            torch.float32: 1e-2,
            torch.half: 1e-2,
            torch.bfloat16: 1e-2,
        }
        for te_output, torch_output in zip(te_outputs[1:], torch_outputs[1:]):
            assert_allclose(te_output, torch_output, atol[dtype], rtol[dtype])

1108

1109
def _test_granular_accuracy(block, bs, dtype, config, delay_wgrad_compute=False, recipe=None):
1110
    reset_rng_states()
1111
1112
1113
    fp8 = recipe is not None
    if fp8:
        FP8GlobalStateManager.reset()
1114
1115

    inp_hidden_states = torch.randn(
1116
        (config.max_seqlen_q, bs, config.hidden_size),
1117
1118
1119
1120
        dtype=dtype,
        device="cuda",
        requires_grad=True,
    )
1121
1122
    inp_hidden_states.retain_grad()

1123
    with autocast(enabled=fp8, recipe=recipe):
1124
1125
1126
        out = block(inp_hidden_states)
        if isinstance(out, (List, Tuple)):
            out = out[0]
1127
1128
    loss = out.sum()
    loss.backward()
1129
1130
    if delay_wgrad_compute:
        block.backward_dw()
1131
1132
1133
1134
1135

    torch.cuda.synchronize()
    outputs = [out, inp_hidden_states.grad]
    for p in block.parameters():
        if p.requires_grad:
1136
1137
1138
1139
1140
            if getattr(p, "main_grad", None) is not None:
                outputs.append(p.main_grad)
                assert p.grad is None  # grad should be None if fuse_wgrad_accumulation is True
            else:
                outputs.append(p.grad)
1141
1142
1143
    return outputs


1144
1145
1146
def _test_dpa_accuracy(block, bs, dtype, config):
    reset_rng_states()

1147
    mask = torch.triu(
1148
1149
        torch.ones(config.max_seqlen_q, config.max_seqlen_kv, dtype=torch.bool, device="cuda"),
        diagonal=1,
1150
    )
1151
    query, key, value = [
1152
        torch.randn(
1153
            (config.max_seqlen_q, bs, config.num_heads, config.kv_channels),
1154
1155
1156
1157
1158
1159
            dtype=dtype,
            device="cuda",
            requires_grad=True,
        )
        for _ in range(3)
    ]
1160
1161
1162
1163
1164

    query.retain_grad()
    key.retain_grad()
    value.retain_grad()

1165
    out = block(query, key, value, attention_mask=mask)
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
    loss = out.sum()
    loss.backward()

    torch.cuda.synchronize()

    return [out, query.grad, key.grad, value.grad]


@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
1176
@pytest.mark.parametrize("model", ["126m"])
1177
1178
1179
1180
1181
def test_dpa_accuracy(dtype, bs, model):
    config = model_configs[model]

    te_dpa = (
        DotProductAttention(
1182
1183
            config.num_heads,
            config.kv_channels,
1184
            attention_dropout=0.0,  # disable dropout, FU uses rng differently
1185
1186
1187
1188
1189
1190
1191
        )
        .to(dtype=dtype)
        .cuda()
    )

    torch_dpa = (
        TorchDotProductAttention(
1192
            config.kv_channels,
1193
            0.0,  # dropout
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
        )
        .to(dtype=dtype)
        .cuda()
    )

    te_outputs = _test_dpa_accuracy(te_dpa, bs, dtype, config)
    torch_outputs = _test_dpa_accuracy(torch_dpa, bs, dtype, config)

    # Check output.
    if dtype == torch.float32:
        assert_allclose(te_outputs[0], torch_outputs[0], 5e-3)
    else:
        assert_allclose(te_outputs[0], torch_outputs[0], 5e-2)

1208
1209
1210
    for te_output, torch_output in zip(te_outputs[1:], torch_outputs[1:]):
        assert_allclose(te_output, torch_output, atol=5e-2, rtol=1e-2)

1211

1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
class TestReturnBiasModule(nn.Module):
    def __init__(self, mod, **kwargs):
        super().__init__()
        self.te_module = mod(**kwargs)
        self.return_bias = kwargs["return_bias"]
        self.bias = kwargs["bias"]

    def forward(self, x):
        if self.return_bias:
            out, bias = self.te_module(x)
            if self.bias:
                out = out + bias
            return out
        return self.te_module(x)


1228
1229
@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
1230
@pytest.mark.parametrize("model", ["small"])
1231
1232
1233
@pytest.mark.parametrize("return_bias", all_boolean)
@pytest.mark.parametrize("bias", all_boolean)
def test_linear_accuracy(dtype, bs, model, return_bias, bias):
1234
1235
    config = model_configs[model]

1236
1237
1238
1239
    te_linear = TestReturnBiasModule(
        Linear,
        in_features=config.hidden_size,
        out_features=4 * config.hidden_size,
1240
        params_dtype=dtype,
1241
1242
        return_bias=return_bias,
        bias=bias,
1243
        device="cuda",
1244
    )
1245

1246
1247
1248
    torch_linear = torch.nn.Linear(
        config.hidden_size,
        4 * config.hidden_size,
1249
        bias=bias,
1250
1251
        device="cuda",
        dtype=dtype,
1252
    )
1253
1254
1255

    # Share params
    with torch.no_grad():
1256
1257
1258
        torch_linear.weight = Parameter(te_linear.te_module.weight.clone())
        if bias:
            torch_linear.bias = Parameter(te_linear.te_module.bias.clone())
1259
1260
1261
1262
1263

    te_outputs = _test_granular_accuracy(te_linear, bs, dtype, config)
    torch_outputs = _test_granular_accuracy(torch_linear, bs, dtype, config)

    # Check output.
1264
1265
1266
1267
1268
1269
1270
1271
1272
    if model == "small":
        tolerance = 5e-3 if dtype == torch.float32 else 5e-2
        rtol = {
            torch.float32: 1.3e-6,
            torch.half: 1e-2,
            torch.bfloat16: 2e-2,
        }
        for te_output, torch_output in zip(te_outputs, torch_outputs):
            assert_allclose(te_output, torch_output, tolerance, rtol[dtype])
1273

1274

1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
@pytest.mark.parametrize("model", ["small"])
@pytest.mark.parametrize("bias", all_boolean)
@pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean)
def test_linear_accuracy_delay_wgrad_compute(dtype, bs, model, bias, fuse_wgrad_accumulation):
    config = model_configs[model]

    te_linear_ref = Linear(
        config.hidden_size,
        4 * config.hidden_size,
        bias=bias,
        params_dtype=dtype,
        device="cuda",
        delay_wgrad_compute=False,
        fuse_wgrad_accumulation=fuse_wgrad_accumulation,
    ).eval()

    te_linear = Linear(
        config.hidden_size,
        4 * config.hidden_size,
        bias=bias,
        params_dtype=dtype,
        device="cuda",
        delay_wgrad_compute=True,
        fuse_wgrad_accumulation=fuse_wgrad_accumulation,
    ).eval()

    # Share params
    with torch.no_grad():
        te_linear_ref.weight = Parameter(te_linear.weight.clone())
        if bias:
            te_linear_ref.bias = Parameter(te_linear.bias.clone())
        if fuse_wgrad_accumulation:
            weight = getattr(te_linear, f"weight")
            weight.main_grad = torch.rand_like(weight, dtype=torch.float32)
            te_linear_ref.weight.main_grad = weight.main_grad.clone()

    te_outputs = _test_granular_accuracy(te_linear, bs, dtype, config, delay_wgrad_compute=True)
    te_outputs_ref = _test_granular_accuracy(
        te_linear_ref, bs, dtype, config, delay_wgrad_compute=False
    )

1318
1319
    # Should be bit-wise match
    for _, (o, o_ref) in enumerate(zip(te_outputs, te_outputs_ref)):
1320
1321
1322
        torch.testing.assert_close(o, o_ref, rtol=0, atol=0)


1323
1324
1325
1326
1327
1328
1329
1330
@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("model", ["small"])
@pytest.mark.parametrize("recipe", fp8_recipes + [None])
def test_linear_accuracy_save_original_input(dtype, model, recipe):
    bs = 1
    fuse_wgrad_accumulation = True
    fp8_model_params = False
    fp8 = recipe is not None
1331
1332
    if fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)
1333
1334
1335
1336
    if fp8 and recipe.delayed():
        pytest.skip("DelayedScaling recipe is not supported with save_original_input")

    config = model_configs[model]
1337
    if config.max_seqlen_q % 16 != 0 and fp8:
1338
1339
        pytest.skip("FP8 requires sequence length to be divisible by 16.")

1340
1341
1342
1343
1344
1345
    if recipe is not None and recipe.nvfp4():
        if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype):
            pytest.skip(
                f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}"
            )

1346
    with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe):
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
        te_linear_ref = Linear(
            config.hidden_size,
            4 * config.hidden_size,
            bias=False,
            params_dtype=dtype,
            device="cuda",
            fuse_wgrad_accumulation=fuse_wgrad_accumulation,
            save_original_input=False,
        ).eval()

        te_linear = Linear(
            config.hidden_size,
            4 * config.hidden_size,
            bias=False,
            params_dtype=dtype,
            device="cuda",
            fuse_wgrad_accumulation=fuse_wgrad_accumulation,
            save_original_input=True,
        ).eval()

    # Share params
    with torch.no_grad():
        te_linear_ref.weight = Parameter(te_linear.weight.clone())
        if fuse_wgrad_accumulation:
            weight = getattr(te_linear, f"weight")
            weight.main_grad = torch.rand_like(weight, dtype=torch.float32)
            te_linear_ref.weight.main_grad = weight.main_grad.clone()

    te_outputs = _test_granular_accuracy(te_linear, bs, dtype, config, recipe=recipe)
    te_outputs_ref = _test_granular_accuracy(te_linear_ref, bs, dtype, config, recipe=recipe)

    # Shoule be bit-wise match
    for i, (o, o_ref) in enumerate(zip(te_outputs, te_outputs_ref)):
        torch.testing.assert_close(o, o_ref, rtol=0, atol=0)


1383
1384
@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
1385
@pytest.mark.parametrize("model", ["126m"])
1386
@pytest.mark.parametrize("eps", [1e-1, 1e-3, 1e-5, 1e-7])
1387
1388
@pytest.mark.parametrize("zero_centered_gamma", all_boolean)
def test_rmsnorm_accuracy(dtype, bs, model, eps, zero_centered_gamma):
1389
1390
    config = model_configs[model]

1391
1392
1393
1394
1395
1396
1397
    te_rmsnorm = RMSNorm(
        config.hidden_size,
        eps=eps,
        params_dtype=dtype,
        zero_centered_gamma=zero_centered_gamma,
        device="cuda",
    ).eval()
1398
1399

    torch_rmsnorm = (
1400
        TorchRMSNorm(config.hidden_size, eps=eps, zero_centered_gamma=zero_centered_gamma)
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
        .to(dtype=dtype)
        .cuda()
        .eval()
    )

    # Share params
    with torch.no_grad():
        torch_rmsnorm.weight = Parameter(te_rmsnorm.weight.clone())

    te_outputs = _test_granular_accuracy(te_rmsnorm, bs, dtype, config)
    torch_outputs = _test_granular_accuracy(torch_rmsnorm, bs, dtype, config)

1413
1414
1415
1416
    atol = {
        torch.float32: 1e-7,
        torch.half: 2e-3,
        torch.bfloat16: 2e-2,
1417
    }
1418
1419

    # Check output.
1420
1421
    assert_allclose(te_outputs[0], torch_outputs[0], atol[dtype])

1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
    atol[torch.float32] = 2e-3
    rtol = {
        torch.float32: 1.3e-6,
        torch.half: 1e-3,
        torch.bfloat16: 1.6e-2,
    }
    # Check gradients
    for te_output, torch_output in zip(te_outputs[1:], torch_outputs[1:]):
        assert_allclose(te_output, torch_output, atol[dtype], rtol[dtype])

1432

1433
1434
@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
1435
@pytest.mark.parametrize("model", ["126m"])
1436
1437
1438
1439
1440
@pytest.mark.parametrize("eps", [1e-1, 1e-3, 1e-5, 1e-7])
@pytest.mark.parametrize("zero_centered_gamma", all_boolean)
def test_layernorm_accuracy(dtype, bs, model, eps, zero_centered_gamma):
    config = model_configs[model]

1441
1442
1443
1444
1445
1446
1447
    te_layernorm = LayerNorm(
        config.hidden_size,
        eps=eps,
        params_dtype=dtype,
        zero_centered_gamma=zero_centered_gamma,
        device="cuda",
    ).eval()
1448
1449

    torch_layernorm = (
1450
        TorchLayerNorm(config.hidden_size, eps=eps, zero_centered_gamma=zero_centered_gamma)
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
        .to(dtype=dtype)
        .cuda()
        .eval()
    )

    # Share params
    with torch.no_grad():
        torch_layernorm.weight = Parameter(te_layernorm.weight.clone())
        torch_layernorm.bias = Parameter(te_layernorm.bias.clone())

    te_outputs = _test_granular_accuracy(te_layernorm, bs, dtype, config)
    torch_outputs = _test_granular_accuracy(torch_layernorm, bs, dtype, config)

1464
1465
1466
1467
    atol = {
        torch.float32: 1e-7,
        torch.half: 2e-3,
        torch.bfloat16: 2e-2,
1468
    }
1469
1470

    # Check output.
1471
    assert_allclose(te_outputs[0], torch_outputs[0], atol[dtype])
1472

1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
    rtol = {
        torch.float32: 1.3e-6,
        torch.half: 1e-3,
        torch.bfloat16: 1.6e-2,
    }
    atol[torch.float32] = 1e-4
    # Check gradients
    for te_output, torch_output in zip(te_outputs[1:], torch_outputs[1:]):
        assert_allclose(te_output, torch_output, atol[dtype], rtol[dtype])

1483

1484
1485
@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
1486
@pytest.mark.parametrize("model", ["small"])
1487
@pytest.mark.parametrize("normalization", all_normalizations)
1488
@pytest.mark.parametrize("zero_centered_gamma", all_boolean)
1489
1490
1491
1492
1493
@pytest.mark.parametrize("return_bias", all_boolean)
@pytest.mark.parametrize("bias", all_boolean)
def test_layernorm_linear_accuracy(
    dtype, bs, model, normalization, zero_centered_gamma, return_bias, bias
):
1494
1495
    config = model_configs[model]

1496
1497
1498
1499
1500
    te_ln_linear = TestReturnBiasModule(
        LayerNormLinear,
        in_features=config.hidden_size,
        out_features=4 * config.hidden_size,
        eps=config.eps,
1501
1502
1503
        normalization=normalization,
        params_dtype=dtype,
        zero_centered_gamma=zero_centered_gamma,
1504
1505
        return_bias=return_bias,
        bias=bias,
1506
        device="cuda",
1507
    )
1508
1509
1510
1511
1512
1513

    torch_ln_linear = (
        TorchLayerNormLinear(
            config.hidden_size,
            4 * config.hidden_size,
            config.eps,
1514
            normalization=normalization,
1515
            zero_centered_gamma=zero_centered_gamma,
1516
            bias=bias,
1517
1518
1519
1520
1521
1522
1523
        )
        .to(dtype=dtype)
        .cuda()
    )

    # Share params
    with torch.no_grad():
1524
1525
1526
        torch_ln_linear.layernorm.weight = Parameter(
            te_ln_linear.te_module.layer_norm_weight.clone()
        )
1527
        if normalization != "RMSNorm":
1528
1529
1530
1531
1532
1533
            torch_ln_linear.layernorm.bias = Parameter(
                te_ln_linear.te_module.layer_norm_bias.clone()
            )
        torch_ln_linear.linear.weight = Parameter(te_ln_linear.te_module.weight.clone())
        if bias:
            torch_ln_linear.linear.bias = Parameter(te_ln_linear.te_module.bias.clone())
1534
1535
1536
1537

    te_outputs = _test_granular_accuracy(te_ln_linear, bs, dtype, config)
    torch_outputs = _test_granular_accuracy(torch_ln_linear, bs, dtype, config)

1538
1539
1540
1541
    atol = {
        torch.float32: 2.5e-4,
        torch.half: 2e-3,
        torch.bfloat16: 2e-2,
1542
    }
1543
1544
1545
1546
1547
    rtol = {
        torch.float32: 1e-3,
        torch.half: 4e-2,
        torch.bfloat16: 4e-2,
    }
1548
1549

    # Check output.
1550
    assert_allclose(te_outputs[0], torch_outputs[0], atol[dtype], rtol[dtype])
1551

1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
    if model == "small":
        atol = {
            torch.float32: 1e-3,
            torch.half: 5e-2,
            torch.bfloat16: 5e-2,
        }
        rtol = {
            torch.float32: 1e-3,
            torch.half: 4e-2,
            torch.bfloat16: 4e-2,
        }
        # Check gradients
        for te_output, torch_output in zip(te_outputs[1:], torch_outputs[1:]):
            assert_allclose(te_output, torch_output, atol[dtype], rtol[dtype])

1567

1568
1569
1570
1571
1572
1573
1574
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
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
@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
@pytest.mark.parametrize("model", ["small"])
@pytest.mark.parametrize("normalization", all_normalizations)
@pytest.mark.parametrize("zero_centered_gamma", all_boolean)
@pytest.mark.parametrize("bias", all_boolean)
@pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean)
def test_layernorm_linear_accuracy_delay_wgrad_compute(
    dtype, bs, model, normalization, zero_centered_gamma, bias, fuse_wgrad_accumulation
):
    config = model_configs[model]

    ln_linear_ref = LayerNormLinear(
        config.hidden_size,
        4 * config.hidden_size,
        config.eps,
        bias=bias,
        normalization=normalization,
        params_dtype=dtype,
        zero_centered_gamma=zero_centered_gamma,
        device="cuda",
        delay_wgrad_compute=False,
        fuse_wgrad_accumulation=fuse_wgrad_accumulation,
    ).eval()

    ln_linear = LayerNormLinear(
        config.hidden_size,
        4 * config.hidden_size,
        config.eps,
        bias=bias,
        normalization=normalization,
        params_dtype=dtype,
        zero_centered_gamma=zero_centered_gamma,
        device="cuda",
        delay_wgrad_compute=True,
        fuse_wgrad_accumulation=fuse_wgrad_accumulation,
    ).eval()

    # Share params
    with torch.no_grad():
        ln_linear_ref.layer_norm_weight = Parameter(ln_linear.layer_norm_weight.clone())
        if normalization != "RMSNorm":
            ln_linear_ref.layer_norm_bias = Parameter(ln_linear.layer_norm_bias.clone())
        ln_linear_ref.weight = Parameter(ln_linear.weight.clone())
        if bias:
            ln_linear_ref.bias = Parameter(ln_linear.bias.clone())
        if fuse_wgrad_accumulation:
            weight = getattr(ln_linear, f"weight")
            weight.main_grad = torch.rand_like(weight, dtype=torch.float32)
            ln_linear_ref.weight.main_grad = weight.main_grad.clone()

    te_outputs = _test_granular_accuracy(ln_linear, bs, dtype, config, delay_wgrad_compute=True)
    te_outputs_ref = _test_granular_accuracy(
        ln_linear_ref, bs, dtype, config, delay_wgrad_compute=False
    )

    # Shoule be bit-wise match
    for i, (o, o_ref) in enumerate(zip(te_outputs, te_outputs_ref)):
        torch.testing.assert_close(o, o_ref, rtol=0, atol=0)


1629
1630
@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
1631
@pytest.mark.parametrize("model", ["small"])
1632
@pytest.mark.parametrize("activation", all_activations)
1633
@pytest.mark.parametrize("normalization", all_normalizations)
1634
1635
1636
@pytest.mark.parametrize("return_bias", all_boolean)
@pytest.mark.parametrize("bias", all_boolean)
def test_layernorm_mlp_accuracy(dtype, bs, model, activation, normalization, return_bias, bias):
1637
1638
    config = model_configs[model]

1639
1640
1641
1642
    te_ln_mlp = TestReturnBiasModule(
        LayerNormMLP,
        hidden_size=config.hidden_size,
        ffn_hidden_size=4 * config.hidden_size,
1643
1644
1645
        activation=activation,
        normalization=normalization,
        params_dtype=dtype,
1646
1647
        return_bias=return_bias,
        bias=bias,
1648
        device="cuda",
1649
    )
1650
1651
1652
1653
1654

    torch_ln_mlp = (
        TorchLayerNormMLP(
            config.hidden_size,
            4 * config.hidden_size,
1655
            activation=activation,
1656
            normalization=normalization,
1657
            bias=bias,
1658
1659
1660
1661
1662
1663
1664
        )
        .to(dtype=dtype)
        .cuda()
    )

    # Share params
    with torch.no_grad():
1665
        torch_ln_mlp.ln.weight = Parameter(te_ln_mlp.te_module.layer_norm_weight.clone())
1666
        if normalization != "RMSNorm":
1667
1668
1669
1670
1671
1672
            torch_ln_mlp.ln.bias = Parameter(te_ln_mlp.te_module.layer_norm_bias.clone())
        torch_ln_mlp.fc1.weight = Parameter(te_ln_mlp.te_module.fc1_weight.clone())
        torch_ln_mlp.fc2.weight = Parameter(te_ln_mlp.te_module.fc2_weight.clone())
        if bias:
            torch_ln_mlp.fc1.bias = Parameter(te_ln_mlp.te_module.fc1_bias.clone())
            torch_ln_mlp.fc2.bias = Parameter(te_ln_mlp.te_module.fc2_bias.clone())
1673
1674
1675
1676

    te_outputs = _test_granular_accuracy(te_ln_mlp, bs, dtype, config)
    torch_outputs = _test_granular_accuracy(torch_ln_mlp, bs, dtype, config)

1677
1678
1679
1680
1681
1682
    atol = {
        torch.float32: 2e-2,
        torch.half: 5e-2,
        torch.bfloat16: 5e-2,
    }

1683
1684
1685
1686
1687
1688
    rtol = {
        torch.float32: 1e-3,
        torch.half: 4e-2,
        torch.bfloat16: 4e-2,
    }

1689
    # Check output.
1690
    assert_allclose(te_outputs[0], torch_outputs[0], atol[dtype], rtol[dtype])
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702

    # Check gradients, only for small model
    rtol = {
        torch.float32: 1e-3,
        torch.half: 1e-2,
        torch.bfloat16: 4e-2,
    }
    atol[torch.half] = 2e-1
    atol[torch.bfloat16] = 2e-1
    if model == "small":
        for te_output, torch_output in zip(te_outputs[1:], torch_outputs[1:]):
            assert_allclose(te_output, torch_output, atol[dtype], rtol[dtype])
1703
1704


1705
@pytest.mark.parametrize("dtype", param_types)
1706
@pytest.mark.parametrize("bs", [2])
1707
1708
1709
1710
@pytest.mark.parametrize("model", ["small"])
@pytest.mark.parametrize("bias", all_boolean)
@pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean)
def test_layernorm_mlp_accuracy_delay_wgrad_compute(
1711
    dtype, bs, model, bias, fuse_wgrad_accumulation
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
):
    config = model_configs[model]

    ln_mlp = LayerNormMLP(
        hidden_size=config.hidden_size,
        ffn_hidden_size=4 * config.hidden_size,
        eps=config.eps,
        bias=bias,
        params_dtype=dtype,
        device="cuda",
        delay_wgrad_compute=True,
        fuse_wgrad_accumulation=fuse_wgrad_accumulation,
    ).eval()

    ln_mlp_ref = LayerNormMLP(
        hidden_size=config.hidden_size,
        ffn_hidden_size=4 * config.hidden_size,
        eps=config.eps,
        bias=bias,
        params_dtype=dtype,
        device="cuda",
        delay_wgrad_compute=False,
        fuse_wgrad_accumulation=fuse_wgrad_accumulation,
    ).eval()

    # Share params
    with torch.no_grad():
        ln_mlp_ref.layer_norm_weight = Parameter(ln_mlp.layer_norm_weight.clone())
1740
        ln_mlp_ref.layer_norm_bias = Parameter(ln_mlp.layer_norm_bias.clone())
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
        ln_mlp_ref.fc1_weight = Parameter(ln_mlp.fc1_weight.clone())
        ln_mlp_ref.fc2_weight = Parameter(ln_mlp.fc2_weight.clone())
        if bias:
            ln_mlp_ref.fc1_bias = Parameter(ln_mlp.fc1_bias.clone())
            ln_mlp_ref.fc2_bias = Parameter(ln_mlp.fc2_bias.clone())
        if fuse_wgrad_accumulation:
            ln_mlp.fc1_weight.main_grad = torch.rand_like(ln_mlp.fc1_weight, dtype=torch.float32)
            ln_mlp_ref.fc1_weight.main_grad = ln_mlp.fc1_weight.main_grad.clone()
            ln_mlp.fc2_weight.main_grad = torch.rand_like(ln_mlp.fc2_weight, dtype=torch.float32)
            ln_mlp_ref.fc2_weight.main_grad = ln_mlp.fc2_weight.main_grad.clone()

    te_outputs = _test_granular_accuracy(ln_mlp, bs, dtype, config, delay_wgrad_compute=True)
    te_outputs_ref = _test_granular_accuracy(
        ln_mlp_ref, bs, dtype, config, delay_wgrad_compute=False
    )

    # Shoule be bit-wise match
    for i, (o, o_ref) in enumerate(zip(te_outputs, te_outputs_ref)):
        torch.testing.assert_close(o, o_ref, rtol=0, atol=0)


1762
def _test_grouped_linear_accuracy(
1763
1764
1765
1766
1767
1768
1769
1770
1771
    block,
    num_gemms,
    bs,
    dtype,
    config,
    recipe,
    fp8,
    fuse_wgrad_accumulation,
    delay_wgrad_compute=False,
1772
):
1773
1774
1775
1776
1777
    reset_rng_states()
    if fp8:
        FP8GlobalStateManager.reset()

    inp_hidden_states = torch.randn(
1778
        (config.max_seqlen_q, bs, config.hidden_size),
1779
1780
1781
1782
1783
1784
        dtype=dtype,
        device="cuda",
        requires_grad=True,
    )
    inp_hidden_states.retain_grad()

1785
    if num_gemms > 1:
1786
1787
        split_size = 1
        if fp8:
1788
            split_size = 16
1789
1790
            if recipe.mxfp8() or recipe.nvfp4():
                split_size = 32
1791
        m = config.max_seqlen_q // split_size
1792
1793
1794
        dist = torch.sort(torch.randint(0, m, (num_gemms - 2,))).values.tolist()
        dist.append(dist[-1])  # Manually add a zero
        m_splits = torch.tensor(dist + [m]) - torch.tensor([0] + dist)
1795
        m_splits = m_splits * split_size
1796
        assert m_splits.sum() == config.max_seqlen_q and len(m_splits) == num_gemms
1797
    else:
1798
        m_splits = torch.tensor([config.max_seqlen_q])
1799

1800
    with autocast(enabled=fp8, recipe=recipe):
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
        if isinstance(block, GroupedLinear):
            m_splits = m_splits * bs
            out = block(inp_hidden_states, m_splits.tolist())
        else:
            out = torch.cat(
                [
                    block[i](inp)
                    for i, inp in enumerate(torch.split(inp_hidden_states, m_splits.tolist()))
                ]
            )
    loss = out.sum()
    loss.backward()
1813
1814
1815
1816
1817
1818
    if delay_wgrad_compute:
        if isinstance(block, GroupedLinear):
            block.backward_dw()
        else:
            for i in range(num_gemms):
                block[i].backward_dw()
1819
1820
1821
1822
1823

    torch.cuda.synchronize()
    outputs = [out, inp_hidden_states.grad]
    for p in block.parameters():
        if p.requires_grad:
1824
1825
1826
1827
1828
            if getattr(p, "main_grad", None) is not None:
                outputs.append(p.main_grad)
                assert p.grad is None  # grad should be None if fuse_wgrad_accumulation is True
            else:
                outputs.append(p.grad)
1829
1830
1831
    return outputs


1832
@pytest.mark.parametrize("dtype", param_types, ids=str)
1833
1834
@pytest.mark.parametrize("num_gemms", [3, 6])
@pytest.mark.parametrize("bs", batch_sizes)
1835
@pytest.mark.parametrize("model", ["126m"])
1836
@pytest.mark.parametrize("recipe", fp8_recipes + [None])
1837
@pytest.mark.parametrize("fp8_model_params", all_boolean)
1838
@pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean)
1839
1840
@pytest.mark.parametrize("bias", all_boolean)
@pytest.mark.parametrize("delay_wgrad_compute", all_boolean)
1841
def test_grouped_linear_accuracy(
1842
1843
1844
1845
1846
1847
1848
    dtype,
    num_gemms,
    bs,
    model,
    recipe,
    fp8_model_params,
    fuse_wgrad_accumulation,
1849
1850
    bias,
    delay_wgrad_compute,
1851
    parallel_mode=None,
1852
    use_cutlass=False,
1853
):
1854
    fp8 = recipe is not None
1855
1856
    if fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)
1857
    if fp8 and fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED:
1858
        pytest.skip("FP8 parameters are not supported in debug mode.")
1859
1860

    config = model_configs[model]
1861
    if config.max_seqlen_q % 16 != 0 and fp8:
1862
1863
        pytest.skip("FP8 requires sequence length to be divisible by 16.")

1864
1865
1866
1867
1868
1869
    if recipe is not None and recipe.nvfp4():
        if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype):
            pytest.skip(
                f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}"
            )

1870
    with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe):
1871
1872
1873
1874
        grouped_linear = GroupedLinear(
            num_gemms,
            config.hidden_size,
            4 * config.hidden_size,
1875
            bias=bias,
1876
            params_dtype=dtype,
1877
            parallel_mode=parallel_mode,
1878
            device="cuda",
1879
            fuse_wgrad_accumulation=fuse_wgrad_accumulation,
1880
            delay_wgrad_compute=delay_wgrad_compute,
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
            save_original_input=False,
        ).eval()
        sequential_linear = torch.nn.ModuleList(
            [
                Linear(
                    config.hidden_size,
                    4 * config.hidden_size,
                    bias=bias,
                    params_dtype=dtype,
                    parallel_mode=parallel_mode,
                    device="cuda",
                    fuse_wgrad_accumulation=fuse_wgrad_accumulation,
                ).eval()
                for _ in range(num_gemms)
            ]
        )

    # Share params
    with torch.no_grad():
        for i in range(num_gemms):
            sequential_linear[i].weight = Parameter(getattr(grouped_linear, f"weight{i}").clone())
            if bias:
                sequential_linear[i].bias = Parameter(getattr(grouped_linear, f"bias{i}").clone())
            if fuse_wgrad_accumulation:
                weight_i = getattr(grouped_linear, f"weight{i}")
                weight_i.main_grad = torch.rand_like(weight_i, dtype=torch.float32)
                sequential_linear[i].weight.main_grad = weight_i.main_grad.clone()
1908
1909
    if IS_HIP_EXTENSION:
        os.environ["NVTE_FORCE_ROCM_GEMM"] = "1"
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
    outputs_ref = _test_grouped_linear_accuracy(
        sequential_linear,
        num_gemms,
        bs,
        dtype,
        config,
        recipe,
        fp8,
        fuse_wgrad_accumulation,
        delay_wgrad_compute,
    )
    outputs = _test_grouped_linear_accuracy(
        grouped_linear,
        num_gemms,
        bs,
        dtype,
        config,
        recipe,
        fp8,
        fuse_wgrad_accumulation,
        delay_wgrad_compute,
    )
1932
1933
    if IS_HIP_EXTENSION:
        os.environ["NVTE_FORCE_ROCM_GEMM"] = "0"
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
    for o, o_ref in zip(outputs, outputs_ref):
        if use_cutlass:
            torch.testing.assert_close(o, o_ref, rtol=1e-3, atol=1e-3)
        else:
            # cuBLAS implementation should be bit-wise match
            torch.testing.assert_close(o, o_ref, rtol=0, atol=0)


@pytest.mark.skipif(
    torch.cuda.get_device_capability() != (9, 0),
    reason="Only enable CUTLASS grouped gemm on Hopper",
)
@pytest.mark.parametrize("dtype", param_types, ids=str)
@pytest.mark.parametrize("num_gemms", [3, 6])
@pytest.mark.parametrize("bs", batch_sizes)
@pytest.mark.parametrize("model", ["126m"])
@pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean)
@pytest.mark.parametrize("delay_wgrad_compute", all_boolean)
def test_grouped_linear_accuracy_cutlass(
    dtype,
    num_gemms,
    bs,
    model,
    fuse_wgrad_accumulation,
    delay_wgrad_compute,
):
    os.environ["NVTE_USE_CUTLASS_GROUPED_GEMM"] = "1"
    test_grouped_linear_accuracy(
        dtype,
        num_gemms,
        bs,
        model,
        None,
        False,
        fuse_wgrad_accumulation,
        False,
        delay_wgrad_compute,
        None,
        use_cutlass=True,
    )
    os.environ.pop("NVTE_USE_CUTLASS_GROUPED_GEMM", None)
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998


@pytest.mark.parametrize("dtype", param_types, ids=str)
@pytest.mark.parametrize("num_gemms", [3])
@pytest.mark.parametrize("bs", [1])
@pytest.mark.parametrize("model", ["126m"])
@pytest.mark.parametrize("recipe", fp8_recipes + [None])
@pytest.mark.parametrize("fp8_model_params", [False])
@pytest.mark.parametrize("fuse_wgrad_accumulation", [True])
@pytest.mark.parametrize("bias", [False])
@pytest.mark.parametrize("delay_wgrad_compute", [True])
def test_grouped_linear_accuracy_save_original_input(
    dtype,
    num_gemms,
    bs,
    model,
    recipe,
    fp8_model_params,
    fuse_wgrad_accumulation,
    bias,
    delay_wgrad_compute,
    parallel_mode=None,
):
    fp8 = recipe is not None
1999
    if fp8 and fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED:
2000
2001
2002
        pytest.skip("FP8 parameters are not supported in debug mode.")
    if fp8 and recipe.delayed():
        pytest.skip("DelayedScaling recipe is not supported with save_original_input")
2003
2004
    if fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)
2005
2006

    config = model_configs[model]
2007
    if config.max_seqlen_q % 16 != 0 and fp8:
2008
2009
        pytest.skip("FP8 requires sequence length to be divisible by 16.")

2010
2011
2012
2013
2014
2015
    if recipe is not None and recipe.nvfp4():
        if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype):
            pytest.skip(
                f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}"
            )

2016
    with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe):
2017
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
        grouped_linear = GroupedLinear(
            num_gemms,
            config.hidden_size,
            4 * config.hidden_size,
            bias=bias,
            params_dtype=dtype,
            parallel_mode=parallel_mode,
            device="cuda",
            fuse_wgrad_accumulation=fuse_wgrad_accumulation,
            delay_wgrad_compute=delay_wgrad_compute,
            save_original_input=True,
2028
2029
2030
2031
2032
2033
        ).eval()
        sequential_linear = torch.nn.ModuleList(
            [
                Linear(
                    config.hidden_size,
                    4 * config.hidden_size,
2034
                    bias=bias,
2035
                    params_dtype=dtype,
2036
                    parallel_mode=parallel_mode,
2037
                    device="cuda",
2038
                    fuse_wgrad_accumulation=fuse_wgrad_accumulation,
2039
2040
2041
2042
2043
2044
2045
2046
2047
                ).eval()
                for _ in range(num_gemms)
            ]
        )

    # Share params
    with torch.no_grad():
        for i in range(num_gemms):
            sequential_linear[i].weight = Parameter(getattr(grouped_linear, f"weight{i}").clone())
2048
2049
            if bias:
                sequential_linear[i].bias = Parameter(getattr(grouped_linear, f"bias{i}").clone())
2050
2051
2052
2053
            if fuse_wgrad_accumulation:
                weight_i = getattr(grouped_linear, f"weight{i}")
                weight_i.main_grad = torch.rand_like(weight_i, dtype=torch.float32)
                sequential_linear[i].weight.main_grad = weight_i.main_grad.clone()
2054

2055
2056
2057
2058
    # Force the sequential_linear and grouped_linear to use hipblaslt rather than hipblas
    if IS_HIP_EXTENSION:
        os.environ["NVTE_FORCE_ROCM_GEMM"] = "1"
    
2059
    outputs_ref = _test_grouped_linear_accuracy(
2060
2061
2062
2063
2064
2065
2066
2067
2068
        sequential_linear,
        num_gemms,
        bs,
        dtype,
        config,
        recipe,
        fp8,
        fuse_wgrad_accumulation,
        delay_wgrad_compute,
2069
2070
    )
    outputs = _test_grouped_linear_accuracy(
2071
2072
2073
2074
2075
2076
2077
2078
2079
        grouped_linear,
        num_gemms,
        bs,
        dtype,
        config,
        recipe,
        fp8,
        fuse_wgrad_accumulation,
        delay_wgrad_compute,
2080
    )
2081
2082
    if IS_HIP_EXTENSION:
        os.environ["NVTE_FORCE_ROCM_GEMM"] = "0"
2083
2084
2085
2086
2087
2088

    # Shoule be bit-wise match
    for i, (o, o_ref) in enumerate(zip(outputs, outputs_ref)):
        torch.testing.assert_close(o, o_ref, rtol=0, atol=0)


2089
@pytest.mark.parametrize("recipe", fp8_recipes + [None])
2090
def test_grouped_linear_accuracy_single_gemm(recipe):
2091
2092
2093
2094
2095
    """Split the tests to save CI time"""
    test_grouped_linear_accuracy(
        dtype=torch.float32,
        num_gemms=1,
        bs=2,
2096
        model="126m",
2097
        recipe=recipe,
2098
        fp8_model_params=True,
2099
        fuse_wgrad_accumulation=True,
2100
2101
        bias=True,
        delay_wgrad_compute=False,
2102
2103
2104
    )


2105
def _test_padding_grouped_linear_accuracy(block, num_gemms, bs, dtype, config, recipe, fp8=False):
2106
2107

    def _pad_tensor_for_fp8(hidden_states, tokens_per_expert):
2108
        align_size = 16
2109
        if recipe.mxfp8() or recipe.nvfp4():
2110
            align_size = 32
2111
        padded_tokens_per_expert = [
2112
2113
            (num_tokens + align_size - 1) // align_size * align_size
            for num_tokens in tokens_per_expert
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
        ]
        hidden_states = torch.split(hidden_states, tokens_per_expert)
        padded_hidden_states = []
        for hidden_state, actual_num_tokens, padded_num_tokens in zip(
            hidden_states, tokens_per_expert, padded_tokens_per_expert
        ):
            padded_hidden_states.append(hidden_state)
            if padded_num_tokens > actual_num_tokens:
                pad_tensor = torch.zeros(
                    padded_num_tokens - actual_num_tokens,
                    hidden_state.shape[1],
                    dtype=hidden_state.dtype,
                    device=hidden_state.device,
                )
                padded_hidden_states.append(pad_tensor)
        padded_hidden_states = torch.cat(padded_hidden_states, dim=0)
        return padded_hidden_states, padded_tokens_per_expert

    def _unpad_tensor_for_fp8(padded_hidden_states, actual_tokens_per_expert, tokens_per_expert):
        inputmats = torch.split(
            padded_hidden_states.view(-1, padded_hidden_states.shape[-1]), tokens_per_expert
        )
        hidden_states = torch.cat(
            [
                grad_output_mat[: actual_tokens_per_expert[i]]
                for i, grad_output_mat in enumerate(inputmats)
            ],
            dim=0,
        )

        return hidden_states

    def _generate_random_numbers(n, total_sum):
        if n <= 0:
            return []

        # reset seed
        random.seed(seed)

        breaks = sorted(random.sample(range(1, total_sum), n - 1))
        random_numbers = (
            [breaks[0]]
            + [breaks[i] - breaks[i - 1] for i in range(1, n - 1)]
            + [total_sum - breaks[-1]]
        )

        return random_numbers

    reset_rng_states()
    if fp8:
        FP8GlobalStateManager.reset()

    inp_hidden_states = torch.randn(
2167
        (config.max_seqlen_q * bs, config.hidden_size),
2168
2169
2170
2171
2172
2173
        dtype=dtype,
        device="cuda",
        requires_grad=True,
    )
    inp_hidden_states.retain_grad()

2174
    m_splits = _generate_random_numbers(num_gemms, config.max_seqlen_q * bs)
2175

2176
    with autocast(enabled=fp8, recipe=recipe):
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
        if isinstance(block, TorchGroupedLinearWithPadding):
            out = block(inp_hidden_states, m_splits)
        else:
            if fp8:
                padded_inp_hidden_states, padding_m_splits = _pad_tensor_for_fp8(
                    inp_hidden_states, m_splits
                )
                padded_inp_hidden_states = block(padded_inp_hidden_states, padding_m_splits)
                out = _unpad_tensor_for_fp8(padded_inp_hidden_states, m_splits, padding_m_splits)
            else:
                out = block(inp_hidden_states, m_splits)

    loss = out.sum()
    loss.backward()

    torch.cuda.synchronize()
    outputs = [out, inp_hidden_states.grad]
    for p in block.parameters():
        if p.requires_grad:
            outputs.append(p.grad)
    return outputs


@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("num_gemms", [3, 6])
@pytest.mark.parametrize("bs", batch_sizes)
2203
@pytest.mark.parametrize("model", ["126m"])
2204
@pytest.mark.parametrize("fp8", [True])
2205
@pytest.mark.parametrize("recipe", fp8_recipes)
2206
2207
@pytest.mark.parametrize("fp8_model_params", all_boolean)
def test_padding_grouped_linear_accuracy(
2208
2209
2210
2211
2212
2213
2214
2215
2216
    dtype,
    num_gemms,
    bs,
    model,
    fp8,
    recipe,
    fp8_model_params,
    parallel_mode=None,
):
2217
2218
    if fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)
2219
2220
2221
2222
    if fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED:
        pytest.skip("FP8 parameters are not supported in debug mode.")

    config = model_configs[model]
2223
    if config.max_seqlen_q % 16 != 0 and fp8:
2224
2225
        pytest.skip("FP8 requires sequence length to be divisible by 16.")

2226
2227
2228
2229
2230
2231
    if recipe is not None and recipe.nvfp4():
        if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype):
            pytest.skip(
                f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}"
            )

2232
    with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe):
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
        grouped_linear = TorchGroupedLinearWithPadding(
            num_gemms,
            config.hidden_size,
            4 * config.hidden_size,
            bias=False,
            params_dtype=dtype,
            parallel_mode=parallel_mode,
            fp8=fp8,
        ).eval()

2243
    with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe):
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
        ref_grouped_linear = GroupedLinear(
            num_gemms,
            config.hidden_size,
            4 * config.hidden_size,
            bias=False,
            params_dtype=dtype,
            parallel_mode=parallel_mode,
            device="cuda",
            save_original_input=False,
        ).eval()

    # Share params
    with torch.no_grad():
        inner_grouped_linear = grouped_linear.linear_fn
        for i in range(num_gemms):
            setattr(
                ref_grouped_linear,
                f"weight{i}",
                Parameter(getattr(inner_grouped_linear, f"weight{i}").clone()),
            )

    outputs = _test_padding_grouped_linear_accuracy(
        grouped_linear, num_gemms, bs, dtype, config, recipe, fp8
    )
    outputs_ref = _test_padding_grouped_linear_accuracy(
        ref_grouped_linear, num_gemms, bs, dtype, config, recipe, fp8
    )

    # Shoule be bit-wise match
    for i, (o, o_ref) in enumerate(zip(outputs, outputs_ref)):
        torch.testing.assert_close(o, o_ref, rtol=0, atol=0)


@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("num_gemms", [3])
@pytest.mark.parametrize("bs", [1])
@pytest.mark.parametrize("model", ["126m"])
@pytest.mark.parametrize("fp8", [True])
@pytest.mark.parametrize("recipe", fp8_recipes)
@pytest.mark.parametrize("fp8_model_params", [False])
def test_padding_grouped_linear_accuracy_save_original_input(
    dtype,
    num_gemms,
    bs,
    model,
    fp8,
    recipe,
    fp8_model_params,
    parallel_mode=None,
2293
):
2294
2295
    if fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED:
        pytest.skip("FP8 parameters are not supported in debug mode.")
2296
2297
    if fp8 and recipe.delayed():
        pytest.skip("DelayedScaling recipe is not supported with save_original_input")
2298
2299
    if fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)
2300
2301

    config = model_configs[model]
2302
    if config.max_seqlen_q % 16 != 0 and fp8:
2303
2304
        pytest.skip("FP8 requires sequence length to be divisible by 16.")

2305
2306
2307
2308
2309
2310
    if recipe is not None and recipe.nvfp4():
        if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype):
            pytest.skip(
                f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}"
            )

2311
    with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe):
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
        grouped_linear = TorchGroupedLinearWithPadding(
            num_gemms,
            config.hidden_size,
            4 * config.hidden_size,
            bias=False,
            params_dtype=dtype,
            parallel_mode=parallel_mode,
            fp8=fp8,
        ).eval()

2322
    with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe):
2323
2324
2325
2326
2327
2328
2329
2330
        ref_grouped_linear = GroupedLinear(
            num_gemms,
            config.hidden_size,
            4 * config.hidden_size,
            bias=False,
            params_dtype=dtype,
            parallel_mode=parallel_mode,
            device="cuda",
2331
            save_original_input=True,
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
        ).eval()

    # Share params
    with torch.no_grad():
        inner_grouped_linear = grouped_linear.linear_fn
        for i in range(num_gemms):
            setattr(
                ref_grouped_linear,
                f"weight{i}",
                Parameter(getattr(inner_grouped_linear, f"weight{i}").clone()),
            )

    outputs = _test_padding_grouped_linear_accuracy(
2345
        grouped_linear, num_gemms, bs, dtype, config, recipe, fp8
2346
2347
    )
    outputs_ref = _test_padding_grouped_linear_accuracy(
2348
        ref_grouped_linear, num_gemms, bs, dtype, config, recipe, fp8
2349
2350
2351
2352
2353
2354
2355
    )

    # Shoule be bit-wise match
    for i, (o, o_ref) in enumerate(zip(outputs, outputs_ref)):
        torch.testing.assert_close(o, o_ref, rtol=0, atol=0)


2356
2357
2358
2359
2360
2361
2362
def _test_gpt_e2e_cuda_graph(block, bs, dtype, config, graph):
    reset_rng_states()

    # Initialize loss function and optimizer.
    loss_fn = torch.nn.MSELoss()
    optimizer = torch.optim.SGD(block.parameters(), lr=0.1)

2363
    # Placeholders used for graph capture.
2364
    static_input = torch.randn(
2365
2366
2367
2368
        config.max_seqlen_q, bs, config.hidden_size, device="cuda", dtype=dtype, requires_grad=True
    )
    static_target = torch.randn(
        config.max_seqlen_q, bs, config.hidden_size, device="cuda", dtype=dtype
2369
    )
2370
2371
2372
2373

    real_input = torch.rand_like(static_input)
    real_target = torch.rand_like(static_target)

2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
    # Basic training loop.
    def train_step():
        optimizer.zero_grad(set_to_none=False)
        out = block(static_input)
        loss = loss_fn(out, static_target)
        loss.backward()
        optimizer.step()
        return out

    # Warmup steps in a separate stream.
    s = torch.cuda.Stream()
    s.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(s):
        for _ in range(3):
            train_step()
    torch.cuda.current_stream().wait_stream(s)

    # Capture graph.
    g = None
    static_output = None
2394
2395
2396
    if graph:
        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g):
2397
2398
2399
2400
2401
2402
2403
            static_output = train_step()

    # Run with new data.
    with torch.no_grad():
        static_input.copy_(real_input)
        static_target.copy_(real_target)
    if graph:
2404
2405
        g.replay()
    else:
2406
        static_output = train_step()
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419

    grads = [static_input.grad]
    for p in block.parameters():
        if p.requires_grad:
            grads.append(p.grad)

    with torch.no_grad():
        output = static_output.clone()
    return output, grads


@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
2420
@pytest.mark.parametrize("model", ["126m"])
2421
def test_gpt_cuda_graph(dtype, bs, model):
2422
2423
    if NVTE_TEST_NVINSPECT_ENABLED:
        pytest.skip("Cuda Graphs are not supported in debug mode.")
2424
2425
2426
2427
2428
2429
    config = model_configs[model]

    sigma = 0.023
    init_method = init_method_normal(sigma)
    output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)

2430
    block_args = (
2431
2432
        config.hidden_size,
        4 * config.hidden_size,
2433
        config.num_heads,
2434
2435
    )
    block_kwargs = dict(
2436
2437
2438
2439
2440
        layernorm_epsilon=config.eps,
        init_method=init_method,
        output_layer_init_method=output_layer_init_method,
        hidden_dropout=0.1,
        attention_dropout=0.1,
2441
        kv_channels=config.kv_channels,
2442
2443
2444
2445
        params_dtype=dtype,
        apply_residual_connection_post_layernorm=False,
        output_layernorm=False,
        device="cuda",
2446
    )
2447
2448
2449
2450
2451
    block = TransformerLayer(*block_args, **block_kwargs)
    graphed_block = TransformerLayer(*block_args, **block_kwargs)
    with torch.no_grad():
        for param1, param2 in zip(block.parameters(), graphed_block.parameters()):
            param2.copy_(param1)
2452

2453
2454
2455
2456
    out, grads = _test_gpt_e2e_cuda_graph(block, bs, dtype, config, False)
    graphed_out, graphed_grads = _test_gpt_e2e_cuda_graph(graphed_block, bs, dtype, config, True)
    params = list(block.parameters())
    graphed_params = list(graphed_block.parameters())
2457

2458
2459
2460
2461
    # Check that results match
    assert_allclose(out, graphed_out, 1e-3)
    assert_allclose(params, graphed_params, 1e-3)
    assert_allclose(grads, graphed_grads, 1e-3)
2462
2463


2464
def _test_gpt_fp8_parameters(bs, dtype, config, fp8_model_params, recipe):
2465
2466
2467
2468
2469
2470
2471
    reset_rng_states()
    FP8GlobalStateManager.reset()

    sigma = 0.023
    init_method = init_method_normal(sigma)
    output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)

2472
    with quantized_model_init(enabled=fp8_model_params, recipe=recipe):
2473
2474
2475
        block = TransformerLayer(
            config.hidden_size,
            4 * config.hidden_size,
2476
            config.num_heads,
2477
2478
2479
2480
2481
            layernorm_epsilon=config.eps,
            init_method=init_method,
            output_layer_init_method=output_layer_init_method,
            hidden_dropout=0.1,
            attention_dropout=0.1,
2482
            kv_channels=config.kv_channels,
2483
2484
2485
2486
2487
            apply_residual_connection_post_layernorm=False,
            output_layernorm=False,
            params_dtype=dtype,
            fuse_qkv_params=True,
            device="cuda",
2488
2489
2490
        )

    te_inp_hidden_states = torch.randn(
2491
        (config.max_seqlen_q, bs, config.hidden_size),
2492
2493
2494
2495
        dtype=dtype,
        device="cuda",
        requires_grad=True,
    )
2496
    te_inp_hidden_states.retain_grad()
2497
    te_inp_attn_mask = get_causal_attn_mask(config.max_seqlen_q)
2498

2499
    with autocast(enabled=True, recipe=recipe):
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
        te_out = block(te_inp_hidden_states, attention_mask=te_inp_attn_mask)
    loss = te_out.sum()
    loss.backward()
    torch.cuda.synchronize()

    outputs = [te_out, te_inp_hidden_states.grad]
    for p in block.parameters():
        if p.requires_grad:
            outputs.append(p.grad)
    return outputs


@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
2514
@pytest.mark.parametrize("model", ["126m"])
2515
2516
@pytest.mark.parametrize("recipe", fp8_recipes)
def test_gpt_fp8_parameters(dtype, bs, model, recipe):
2517
2518
    if not fp8_available:
        pytest.skip(reason_for_no_fp8)
2519
2520
    if NVTE_TEST_NVINSPECT_ENABLED:
        pytest.skip("FP8 parameters are not supported in debug mode.")
2521

2522
2523
2524
2525
2526
2527
    if recipe.nvfp4():
        if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype):
            pytest.skip(
                f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}"
            )

2528
2529
    config = model_configs[model]

2530
2531
    outputs = _test_gpt_fp8_parameters(bs, dtype, config, False, recipe)
    outputs_fp8_params = _test_gpt_fp8_parameters(bs, dtype, config, True, recipe)
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543

    # Check that results match
    tols = dict(rtol=0.125, atol=0.0675)
    for i, (ref, test) in enumerate(zip(outputs, outputs_fp8_params)):
        torch.testing.assert_close(
            test,
            ref,
            msg=f"Mismatch in tensor {i}",
            rtol=0.125,
            atol=0.0675,
        )

2544
2545
2546

@pytest.mark.parametrize("dtype", param_types)
@pytest.mark.parametrize("bs", batch_sizes)
2547
@pytest.mark.parametrize("model", ["126m"])
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
def test_transformer_layer_hidden_states_format(dtype, bs, model):
    config = model_configs[model]

    sigma = 0.023
    init_method = init_method_normal(sigma)
    output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)

    # Set `torch.manual_seed` to make sure the weights are identical to the
    # other layer. Set `*dropout` values to 0 to make sure the forward pass
    # is identical to the other layer.
    torch.manual_seed(0)
2559
2560
2561
    block_sbhd = TransformerLayer(
        config.hidden_size,
        4 * config.hidden_size,
2562
        config.num_heads,
2563
2564
2565
2566
2567
        layernorm_epsilon=config.eps,
        init_method=init_method,
        output_layer_init_method=output_layer_init_method,
        hidden_dropout=0,
        attention_dropout=0,
2568
        kv_channels=config.kv_channels,
2569
2570
2571
2572
2573
        params_dtype=dtype,
        apply_residual_connection_post_layernorm=False,
        output_layernorm=False,
        device="cuda",
        attn_input_format="sbhd",
2574
2575
2576
2577
2578
2579
    )

    # Set `torch.manual_seed` to make sure the weights are identical to the
    # other layer. Set `*dropout` values to 0 to make sure the forward pass
    # is identical to the other layer.
    torch.manual_seed(0)
2580
2581
2582
    block_bshd = TransformerLayer(
        config.hidden_size,
        4 * config.hidden_size,
2583
        config.num_heads,
2584
2585
2586
2587
2588
        layernorm_epsilon=config.eps,
        init_method=init_method,
        output_layer_init_method=output_layer_init_method,
        hidden_dropout=0,
        attention_dropout=0,
2589
        kv_channels=config.kv_channels,
2590
2591
2592
2593
2594
        params_dtype=dtype,
        apply_residual_connection_post_layernorm=False,
        output_layernorm=False,
        device="cuda",
        attn_input_format="bshd",
2595
2596
    )

2597
2598
2599
2600
    torch.manual_seed(0)
    block_thd = TransformerLayer(
        config.hidden_size,
        4 * config.hidden_size,
2601
        config.num_heads,
2602
2603
2604
2605
2606
        layernorm_epsilon=config.eps,
        init_method=init_method,
        output_layer_init_method=output_layer_init_method,
        hidden_dropout=0,
        attention_dropout=0,
2607
        kv_channels=config.kv_channels,
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
        params_dtype=dtype,
        apply_residual_connection_post_layernorm=False,
        output_layernorm=False,
        device="cuda",
        attn_input_format="thd",
        self_attn_mask_type="padding_causal",
    )

    for (n1, p1), (n2, p2), (n3, p3) in zip(
        block_bshd.named_parameters(), block_sbhd.named_parameters(), block_thd.named_parameters()
    ):
        assert torch.all(torch.eq(p1, p2) & torch.eq(p1, p3)), f"{n1}, {n2} and {n3} not identical"
2620
2621

    x_sbhd = torch.randn(
2622
        (config.max_seqlen_q, bs, config.hidden_size),
2623
2624
2625
2626
        dtype=dtype,
        device="cuda",
        requires_grad=True,
    )
2627

2628
    x_bshd = x_sbhd.transpose(0, 1).contiguous()
2629
2630
    x_thd = x_bshd.reshape(bs * config.max_seqlen_q, config.hidden_size).contiguous()
    x_thd_cumsum = torch.arange(bs + 1, device="cuda", dtype=torch.int32) * config.max_seqlen_q
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641

    # To make sure forward is also identical (just in case some module decides
    # to act fancy)
    torch.manual_seed(0)
    y_sbhd = block_sbhd(x_sbhd)

    # To make sure forward is also identical (just in case some module decides
    # to act fancy)
    torch.manual_seed(0)
    y_bshd = block_bshd(x_bshd)

2642
2643
2644
    # Check that results match
    torch.testing.assert_close(
        y_bshd,
2645
        y_sbhd.transpose(0, 1).contiguous(),
2646
    )
2647

2648
2649
2650
2651
2652
2653
2654
2655
2656
    # THD is not supported in float32 and on GPUs older than Ampere, skip the test here
    if dtype != torch.float32 and sm_80plus:
        # To make sure forward is also identical (just in case some module decides
        # to act fancy)
        torch.manual_seed(0)
        y_thd = block_thd(
            x_thd,
            cu_seqlens_q=x_thd_cumsum,
            cu_seqlens_kv=x_thd_cumsum,
2657
2658
            max_seqlen_q=config.max_seqlen_q,
            max_seqlen_kv=config.max_seqlen_kv,
2659
2660
2661
2662
        )

        torch.testing.assert_close(
            y_bshd,
2663
            y_thd.reshape(bs, config.max_seqlen_q, config.hidden_size).contiguous(),
2664
        )
2665

2666
2667
2668
2669
2670
2671
2672
2673
2674
2675

@pytest.mark.parametrize(
    "shape",
    [
        (1, 127, 128, 512),
        (8, 15, 128, 512),
        (8, 1027, 128, 512),
        (16, 10027, 128, 512),
    ],
)
2676
@pytest.mark.parametrize("dtype", param_types, ids=str)
2677
2678
@pytest.mark.parametrize("layout", ["TN", "NN", "NT"])
@pytest.mark.parametrize("accumulate", [False, True])
2679
2680
@pytest.mark.parametrize("use_cutlass", use_cutlass_grouped_gemm)
def test_grouped_gemm(shape, dtype, layout, accumulate, use_cutlass):
2681
2682
2683
2684
2685
2686
2687
2688
2689
2690
    torch.manual_seed(0)
    z, m, k, n = shape

    dist = torch.sort(torch.randint(0, m, (z - 1,))).values.tolist()
    m_splits = torch.tensor(dist + [m]) - torch.tensor([0] + dist)
    assert m_splits.sum() == m and len(m_splits) == z
    m_splits = m_splits.tolist()

    if layout == "TN":
        A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)]  # weight
2691
2692
2693
        B = list(torch.split(torch.randn(m, k, dtype=dtype, device="cuda"), m_splits))  # input
        out = [torch.randn(m, n, dtype=dtype, device="cuda")]  # output
        out_ref = [o.clone() for o in torch.split(out[0], m_splits)]
2694
        grad = False
2695
        single_output = True
2696
2697
    elif layout == "NN":
        A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)]  # weight
2698
2699
2700
2701
2702
        B = list(
            torch.split(torch.randn(m, n, dtype=dtype, device="cuda"), m_splits)
        )  # grad_output
        out = [torch.randn(m, k, dtype=dtype, device="cuda")]  # dgrad
        out_ref = [o.clone() for o in torch.split(out[0], m_splits)]
2703
        grad = True
2704
        single_output = True
2705
    else:  # layout == "NT"
2706
2707
2708
2709
        A = list(torch.split(torch.randn(m, k, dtype=dtype, device="cuda"), m_splits))  # input
        B = list(
            torch.split(torch.randn(m, n, dtype=dtype, device="cuda"), m_splits)
        )  # grad_output
2710
        out = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)]  # wgrad
2711
        out_ref = [o.clone() for o in out]
2712
        grad = True
2713
        single_output = False
2714

2715
2716
2717
    if use_cutlass:
        os.environ["NVTE_USE_CUTLASS_GROUPED_GEMM"] = "1"

2718
2719
    # Force the sequential_linear and grouped_linear to use hipblaslt rather than hipblas
    if IS_HIP_EXTENSION:
2720
        ori_force_rocm_gemm = os.environ.get("NVTE_FORCE_ROCM_GEMM", None)
2721
2722
        os.environ["NVTE_FORCE_ROCM_GEMM"] = "1"

2723
    for i in range(z):
2724
        general_gemm(
2725
2726
            A[i],
            B[i],
2727
            dtype,
2728
2729
2730
2731
2732
            grad=grad,
            accumulate=accumulate,
            layout=layout,
            out=out_ref[i],
        )
2733
2734
    if single_output:
        out_ref = [torch.cat(out_ref)]
2735

2736
    general_grouped_gemm(
2737
        A,
2738
2739
        B,
        out,
2740
        dtype,
2741
        m_splits=m_splits,
2742
2743
2744
        grad=grad,
        accumulate=accumulate,
        layout=layout,
2745
        single_output=single_output,
2746
    )
2747
    if IS_HIP_EXTENSION:
2748
2749
2750
2751
        if ori_force_rocm_gemm is not None:
            os.environ["NVTE_FORCE_ROCM_GEMM"] = ori_force_rocm_gemm
        else:
            del os.environ["NVTE_FORCE_ROCM_GEMM"]
2752
2753

    for o, o_ref in zip(out, out_ref):
2754
2755
2756
2757
2758
2759
2760
2761
        if not use_cutlass:
            # cublas implementation should be bit-wise match
            torch.testing.assert_close(o, o_ref, rtol=0, atol=0)
        else:
            torch.testing.assert_close(o, o_ref, rtol=1.5e-2, atol=1.5e-2)

    if use_cutlass:
        os.environ.pop("NVTE_USE_CUTLASS_GROUPED_GEMM", None)
2762
2763


2764
2765
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
2795
2796
2797
2798
2799
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
2812
2813
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
@pytest.mark.parametrize("N", [32])
@pytest.mark.parametrize("datatype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize(
    "input_quantizer",
    [
        Float8CurrentScalingQuantizer(fp8_dtype=tex.DType.kFloat8E4M3, device="cuda"),
        MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3),
    ],
)
@pytest.mark.parametrize(
    "out_quantizer",
    [
        Float8CurrentScalingQuantizer(fp8_dtype=tex.DType.kFloat8E4M3, device="cuda"),
        MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3),
        Float8Quantizer(
            torch.ones(1).cuda().squeeze(), torch.ones(1).cuda().squeeze(), tex.DType.kFloat8E4M3
        ),
    ],
)
def test_fp8gemm_with_unfused_quantization(N, datatype, input_quantizer, out_quantizer):
    # For MXFP8 and CurrentScaling, below unfused quantization should happen
    # FP8 input --> cublas GEMM --> BF16 output --> Quantize to FP8 --> fp8 Output
    # Skip invalid configurations
    is_mxfp8_needed = isinstance(input_quantizer, MXFP8Quantizer) or isinstance(
        out_quantizer, MXFP8Quantizer
    )
    if not fp8_available:
        pytest.skip(reason_for_no_fp8)
    if is_mxfp8_needed and not mxfp8_available:
        pytest.skip(reason_for_no_mxfp8)
    inp_fp8 = input_quantizer(torch.randn(N, N, device="cuda", dtype=datatype))
    weight_fp8 = input_quantizer(torch.randn(N, N, device="cuda", dtype=datatype))
    outp_type = torch.float32
    quantized_out, *_ = general_gemm(
        weight_fp8,
        inp_fp8,
        outp_type,
        quantization_params=out_quantizer,
        bias=None,
        use_split_accumulator=False,
    )

    out, *_ = general_gemm(
        weight_fp8,
        inp_fp8,
        outp_type,
        quantization_params=None,
        bias=None,
        use_split_accumulator=False,
    )
    expected_quantized_out = out_quantizer(out)

    # Match results again Pytorch GEMM and allow for quantization tolerance
    pytorch_out = torch.matmul(
        inp_fp8.dequantize().to(torch.float64),
        torch.transpose(weight_fp8.dequantize().to(torch.float64), 0, 1),
    )
    fp8_tols = dict(rtol=0.125, atol=0.0675)
    torch.testing.assert_close(
        pytorch_out.to(outp_type), expected_quantized_out.dequantize(), **fp8_tols
    )
    # Match results between quantization happening inside vs outside general_gemm
    torch.testing.assert_close(expected_quantized_out.dequantize(), quantized_out.dequantize())
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837


@pytest.mark.parametrize(
    "shape",
    [
        (1, 128, 128, 512),
        (8, 1024, 128, 512),
        (16, 4096, 128, 512),
    ],
)
@pytest.mark.parametrize("accumulate", [False, True])
2838
def test_fp8_grouped_gemm(shape, accumulate):
2839
2840
2841
2842
    if not fp8_available:
        pytest.skip(reason_for_no_fp8)

    z, m, k, n = shape
2843
    m_splits = [m // z] * z
2844
2845
2846
2847
2848
2849
2850
2851

    dtype = torch.bfloat16
    A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)]  # weight
    B = torch.split(torch.randn(m, k, dtype=dtype, device="cuda"), m_splits)  # input
    out = torch.split(torch.randn(m, n, dtype=dtype, device="cuda"), m_splits)  # output
    out_ref = [o.clone() for o in out]

    # fp8 should be robust enough to this fake scale
2852
2853
    scale = 1 + torch.rand(1, dtype=torch.float32, device="cuda").squeeze()
    amax = torch.zeros(1, 1, dtype=torch.float32, device="cuda")
2854

2855
2856
2857
2858
    a_quantizers = [
        Float8Quantizer(
            scale.clone(),
            amax.clone(),
2859
2860
            tex.DType.kFloat8E4M3,
        )
2861
        for _ in range(z)
2862
    ]
2863
2864
2865
2866
2867
    b_quantizers = [
        Float8Quantizer(
            scale.clone(),
            amax.clone(),
            tex.DType.kFloat8E4M3,
2868
        )
2869
        for _ in range(z)
2870
2871
    ]

2872
2873
2874
2875
2876
2877
    A_fp8 = []
    B_fp8 = []

    for i in range(z):
        A_fp8.append(a_quantizers[i](A[i]))
        B_fp8.append(b_quantizers[i](B[i]))
2878
2879
2880

    # baseline
    for i in range(z):
2881
        general_gemm(
2882
2883
            A_fp8[i],
            B_fp8[i],
2884
            dtype,
2885
2886
2887
            out=out_ref[i],
            accumulate=accumulate,
        )
2888
2889
2890
2891
2892
    general_grouped_gemm(
        A_fp8,
        B_fp8,
        out,
        dtype,
2893
        m_splits=m_splits,
2894
2895
        accumulate=accumulate,
    )
2896
2897
2898
2899

    # should be bit-wise match
    for o, o_ref in zip(out, out_ref):
        torch.testing.assert_close(o, o_ref, rtol=0, atol=0)
2900
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
2914
2915
2916
2917
2918
2919
2920
2921
2922
2923
2924
2925
2926
2927
2928
2929
2930
2931
2932
2933
2934
2935
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948


def test_noncontiguous():
    def _create2modules(m, params):
        mod1 = m(*params)
        mod2 = m(*params)
        for p1, p2 in zip(mod1.parameters(), mod2.parameters()):
            p2.data = p1.data.clone()

        return mod1, mod2

    def _run_module(m, inp):
        out = m(inp)
        out.sum().backward()
        ret = [out]
        if inp.grad is not None:
            ret.append(inp.grad)

        for p in m.parameters():
            if p.requires_grad:
                ret.append(p.grad)
        return ret

    a = torch.randn((128, 256), device="cuda", requires_grad=True)
    a = a.T
    assert not a.is_contiguous(), "The test is supposed to test noncontiguous input."

    b = a.contiguous()

    # LayerNorm
    ln1, ln2 = _create2modules(LayerNorm, [128])
    outT = _run_module(ln1, a)
    out = _run_module(ln2, b)

    assert_allclose(out, outT, 1e-7)

    # RMSNorm
    ln1, ln2 = _create2modules(RMSNorm, [128])
    outT = _run_module(ln1, a)
    out = _run_module(ln2, b)

    assert_allclose(out, outT, 1e-7)

    # GEMM
    g1, g2 = _create2modules(Linear, [128, 128])
    outT = _run_module(g1, a)
    out = _run_module(g2, b)

    assert_allclose(out, outT, 1e-7)