[PyTorch] Fix Float8Tensor transpose caching in #575

docs/api/pytorch.rst

@@ -41,4 +41,6 @@ pyTorch
 
 .. autoapifunction:: transformer_engine.pytorch.onnx_export
 
+.. autoapifunction:: transformer_engine.pytorch.make_graphed_callables
+
 .. autoapifunction:: transformer_engine.pytorch.get_cpu_offload_context

tests/pytorch/fused_attn/test_fused_attn.py

@@ -33,10 +33,7 @@
     CudaRNGStatesTracker,
 )
 import transformer_engine.pytorch.fp8 as fp8
-from transformer_engine.pytorch.module.base import (
-    TransformerEngineBaseModule,
-    _prepare_backward,
-)
+from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
 from transformer_engine.pytorch.utils import (
     get_device_compute_capability,
     init_method_normal,
@@ -1188,8 +1185,7 @@ def forward(
     def backward(
         ctx, grad_output: torch.Tensor
     ) -> Tuple[Union[torch.Tensor, None], ...]:
-
-        with _prepare_backward(True, ctx.fp8_meta, None, 1, name="_DPA"):
+        with torch.cuda.nvtx.range("_DPA"):
             (
                 inputmat_t,
                 qkv_weight_t_fp8,

tests/pytorch/test_cuda_graphs.py

@@ -0,0 +1,175 @@
+"""Cuda graphs tests."""
+import argparse
+
+import torch
+import transformer_engine.pytorch as te
+import apex
+
+
+def str_to_optimizer(optim):
+    """Get optimizer."""
+    if optim == "sgd":
+        return torch.optim.SGD
+    if optim == "adamw":
+        return torch.optim.AdamW
+    if optim == "fused_sgd":
+        return apex.optimizers.FusedSGD
+    return apex.optimizers.FusedAdam
+
+
+def str_to_torch_dtype(dtype):
+    """Get pytorch dtype."""
+    if dtype == "bf16":
+        return torch.bfloat16
+    if dtype == "fp16":
+        return torch.float16
+    return torch.float32
+
+
+def manual_seed(seed):
+    """Set seed."""
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+
+
+def generate_data(args, warmup=False, gen_labels=False):
+    """Generate synthetic data."""
+    dtype = str_to_torch_dtype(args.dtype)
+    gen_func = torch.ones if warmup else torch.randn
+    if args.module == "dpa":
+        inputs = [gen_func(
+            args.seq_length, args.bs, args.nheads,
+            args.embed, device="cuda", requires_grad=True, dtype=dtype
+        ) for _ in range(3)]
+    else:
+        inputs = [gen_func(args.seq_length, args.bs,
+                              args.hdim, device="cuda", requires_grad=True, dtype=dtype)]
+
+    if not gen_labels:
+        return inputs
+
+    target = torch.randn(args.seq_length, args.bs, args.hdim, device="cuda", dtype=dtype)
+    return inputs, target
+
+
+def print_values(model, output):
+    """Debug."""
+    values = []
+    for param in model.parameters():
+        values.append(param.sum().item())
+        if param.grad is not None:
+            values.append(param.grad.sum().item())
+    values.append(output.sum().item())
+    print(values)
+
+
+def parse_args():
+    """Arguments."""
+    parser = argparse.ArgumentParser(description="Args for testing CUDA graphs with TE layers.")
+    parser.add_argument('--seed', type=int, default=1234)
+    parser.add_argument('--dtype', type=str, default="bf16", choices=["bf16", "fp16", "fp32"])
+    parser.add_argument('--optimizer', type=str, default="adamw",
+                        choices=["fused_adamw", "fused_sgd", "sgd", "adamw"])
+    parser.add_argument('--num-layers', type=int, default=1)
+    parser.add_argument('--module', default="linear",
+                        choices=['linear', 'layernorm_linear', 'layernorm_mlp',
+                                 'transformer', 'dpa', 'mha'])
+    parser.add_argument('--fp8', action='store_true')
+    parser.add_argument('--fp8-params', action='store_true')
+    parser.add_argument('--graph', action='store_true')
+    parser.add_argument('--graph-mode', default="full", choices=['full', 'individual'])
+    parser.add_argument('--num-warmup-iters', type=int, default=3)
+    parser.add_argument('--steps', type=int, default=1)
+    parser.add_argument('--hdim', type=int, default=768)
+    parser.add_argument('--seq-length', type=int, default=2048)
+    parser.add_argument('--bs', type=int, default=2)
+    parser.add_argument('--nheads', type=int, default=12)
+    parser.add_argument('--dropout', type=float, default=0.1)
+    return parser.parse_args()
+
+
+def train(args):
+    """Train."""
+
+    dtype = str_to_torch_dtype(args.dtype)
+    if args.fp8_params:
+        assert args.fp8, "FP8 execution needed for FP8 parameters."
+        assert (args.optimizer in ("sgd", "adamw")
+        ), f"Unsupported optimizer {args.optimizer} for FP8 parameters."
+
+    with te.fp8_model_init(enabled=args.fp8_params):
+        # Create modules.
+        if args.module == "transformer":
+            modules = [te.TransformerLayer(
+                            args.hdim, args.hdim, args.nheads,
+                            hidden_dropout=args.dropout,
+                            attention_dropout=args.dropout,
+                            fuse_qkv_params=True,
+                            params_dtype=dtype,
+                        ) for _ in range(args.num_layers)]
+        elif args.module == "layernorm_mlp":
+            modules = [te.LayerNormMLP(
+                args.hdim, args.hdim, params_dtype=dtype
+            ) for _ in range(args.num_layers)]
+        elif args.module == "layernorm_linear":
+            modules = [te.LayerNormLinear(
+                args.hdim, args.hdim, params_dtype=dtype
+            ) for _ in range(args.num_layers)]
+        elif args.module == "mha":
+            modules = [te.MultiheadAttention(
+                args.hdim, args.nheads, attention_dropout=args.dropout, params_dtype=dtype
+            ) for _ in range(args.num_layers)]
+        elif args.module == "dpa":
+            assert args.hdim % args.nheads == 0, "Err."
+            assert args.num_layers == 1, "Err."
+            args.embed = args.hdim // args.nheads
+            modules = [te.DotProductAttention(
+                        args.nheads, args.embed, attention_dropout=args.dropout
+                        ) for _ in range(args.num_layers)]
+        else:
+            modules = [te.Linear(
+                args.hdim, args.hdim, device="cuda", params_dtype=dtype
+            ) for _ in range(args.num_layers)]
+
+        # Generate model and wrap API to return graphed version.
+        if args.graph:
+            # Graph entire module at once.
+            if args.graph_mode == "full":
+                model = modules[0] if args.module == "dpa" else torch.nn.Sequential(*modules)
+                model = te.make_graphed_callables(
+                        model,
+                        generate_data(args, warmup=True),
+                        num_warmup_iters=args.num_warmup_iters,
+                        enabled=args.fp8)
+            else:
+                modules = [te.make_graphed_callables(
+                    module,
+                    generate_data(args, warmup=True),
+                    num_warmup_iters=args.num_warmup_iters,
+                    enabled=args.fp8) for module in modules]
+                model = modules[0] if args.module == "dpa" else torch.nn.Sequential(*modules)
+        else:
+            model = modules[0] if args.module == "dpa" else torch.nn.Sequential(*modules)
+
+    # Loss function and optimizer.
+    loss_fn = torch.nn.MSELoss()
+    optimizer = str_to_optimizer(args.optimizer)(model.parameters(), lr=0.001)
+
+    # Launch.
+    for _ in range(args.steps):
+        inputs, target = generate_data(args, gen_labels=True)
+        with te.fp8_autocast(enabled=args.fp8):
+            output = model(*inputs)
+        loss = loss_fn(output, target)
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+
+    # Debug.
+    print_values(model, output)
+
+
+if __name__ == "__main__":
+    arguments = parse_args()
+    manual_seed(arguments.seed)
+    train(arguments)

tests/pytorch/test_float8tensor.py

@@ -258,11 +258,9 @@ def test_inplace_ops(
             torch.testing.assert_close(x_fp8, x_ref, **tols)
 
     @pytest.mark.parametrize("dims", [[33, 41], [5, 7, 11]])
-    @pytest.mark.parametrize("transpose_dims", [(0, 1), (-2, -1), (0, 0)])
     def test_transpose(
         self,
         dims: DimsType,
-        transpose_dims: Tuple[int, int],
         fp8_dtype: tex.DType = tex.DType.kFloat8E4M3,
         scale: float = 0.5,
         dtype: torch.dtype = torch.float32,
@@ -280,65 +278,65 @@ def test_transpose(
         x_ref = x_fp8.from_float8()
 
         # Perform transpose
-        y_fp8 = x_fp8.transpose(*transpose_dims)
-        y_ref = x_ref.transpose(*transpose_dims)
+        y_fp8 = Float8Tensor.make_like(x_fp8, data=x_fp8._data_transpose())
+        y_ref = x_ref.reshape(-1, dims[-1]).transpose(0, 1)
 
         # Check results
         tols = dict(rtol=0, atol=0)
         torch.testing.assert_close(y_fp8, y_ref, **tols)
 
         # Make sure we are not trivially passing the test
-        if transpose_dims[0] != transpose_dims[1]:
-            with pytest.raises(AssertionError):
-                torch.testing.assert_close(
-                    y_fp8,
-                    x_ref,
-                    **tols,
-                )
-
-        # Check transpose caching
-        if x_fp8.dim() == 2 and transpose_dims[0] != transpose_dims[1]:
-
-            # Check that cached transpose is returned when expected
-            # Note: Sneakily destroy data so that recalculating
-            # transpose would give wrong answer.
-            x_fp8 += 0.5
-            x_ref = x_fp8.from_float8()
-            torch.testing.assert_close(
-                x_fp8.transpose(*transpose_dims, update_cache="lazy"),
-                x_ref.transpose(*transpose_dims),
-                **tols,
-            )
-            x_fp8_data = x_fp8._data.clone()
-            x_fp8._data.zero_()
-            torch.testing.assert_close(
-                x_fp8.transpose(*transpose_dims),
-                x_ref.transpose(*transpose_dims),
-                **tols,
-            )
-            torch.testing.assert_close(
-                x_fp8.transpose(*transpose_dims, update_cache="lazy"),
-                x_ref.transpose(*transpose_dims),
-                **tols,
-            )
-            torch.testing.assert_close(
-                x_fp8.transpose(*transpose_dims, update_cache="force"),
-                torch.zeros_like(x_ref.transpose(*transpose_dims)),
-                rtol=0,
-                atol=0,
-            )
-            x_fp8._data.copy_(x_fp8_data)
-            x_fp8._reset_caches()
-
-            # Make sure cache is reset after in-place operation
-            x_fp8.transpose(*transpose_dims, update_cache="force")
-            x_fp8 += 0.5
-            x_ref = x_fp8.from_float8()
-            torch.testing.assert_close(
-                x_fp8.transpose(*transpose_dims),
-                x_ref.transpose(*transpose_dims),
-                **tols,
-            )
+        with pytest.raises(AssertionError):
+            torch.testing.assert_close(y_fp8, x_ref, **tols)
+
+        # Check that cached transpose is returned when expected
+        # Note: Sneakily destroy data so that recalculating
+        # transpose would give wrong answer.
+        x_fp8 += 0.5
+        x_ref = x_fp8.from_float8()
+        y_ref = x_ref.reshape(-1, dims[-1]).transpose(0, 1)
+        torch.testing.assert_close(
+            Float8Tensor.make_like(
+                x_fp8,
+                data=x_fp8._data_transpose(fill_cache=True),
+            ),
+            y_ref,
+            **tols,
+        )
+        x_fp8_data = x_fp8._data.clone()
+        x_fp8._data.zero_()
+        torch.testing.assert_close(
+            Float8Tensor.make_like(
+                x_fp8,
+                data=x_fp8._data_transpose(),
+            ),
+            y_ref,
+            **tols,
+        )
+        torch.testing.assert_close(
+            Float8Tensor.make_like(
+                x_fp8,
+                data=x_fp8._data_transpose(force_compute=True),
+            ),
+            torch.zeros_like(y_ref),
+            rtol=0,
+            atol=0,
+        )
+        x_fp8._data.copy_(x_fp8_data)
+        x_fp8._reset_caches()
+
+        # Make sure cache is reset after in-place operation
+        x_fp8._data_transpose(fill_cache=True)
+        x_fp8 += 0.5
+        x_ref = x_fp8.from_float8()
+        torch.testing.assert_close(
+            Float8Tensor.make_like(
+                x_fp8,
+                data=x_fp8._data_transpose(),
+            ),
+            x_ref.reshape(-1, dims[-1]).transpose(0, 1),
+            **tols,
+        )
 
     def test_serialization(
         self,