Skip to content

GitLab

"...AutoBuildImmortalWrt.git" did not exist on "d91208ec7a4313b67daffab477e4407c8f83702f"
Unverified Commit 49a4535d authored by Jaemin Choi's avatar Jaemin Choi Committed by GitHub
Browse files

Add NVTX ranges to categorize execution (#1447) 


Signed-off-by: default avatarJaemin Choi <jaeminc@nvidia.com>
Signed-off-by: default avatarTim Moon <tmoon@nvidia.com>
Co-authored-by: default avatarJaemin Choi <jaeminc@nvidia.com>
Co-authored-by: default avatarTim Moon <tmoon@nvidia.com>
parent b87e539d
Hide whitespace changes
Inline Side-by-side
Showing with 150 additions and 6 deletions
transformer_engine/pytorch/attention.py
View file @ 49a4535d
......@@ -23,7 +23,11 @@ import torch.nn.functional as F
import transformer_engine_torch as tex
import transformer_engine as te
from transformer_engine.pytorch.utils import get_cudnn_version
from transformer_engine.pytorch.utils import (
get_cudnn_version,
nvtx_range_pop,
nvtx_range_push,
)
from transformer_engine.pytorch.cpp_extensions.fused_attn import (
fused_attn_fwd,
fused_attn_bwd,
......@@ -1834,6 +1838,7 @@ class AttnFuncWithCPAndKVP2P(torch.autograd.Function):
quantizers,
):
# pylint: disable=missing-function-docstring
nvtx_range_push("transformer_engine.AttnFuncWithCPAndKVP2P.forward")
if softmax_scale is None:
softmax_scale = q.shape[-1] ** (-0.5)
......@@ -2756,12 +2761,14 @@ class AttnFuncWithCPAndKVP2P(torch.autograd.Function):
ctx.fp8_meta = fp8_meta
ctx.is_input_fp8 = is_input_fp8
ctx.is_output_fp8 = is_output_fp8
nvtx_range_pop("transformer_engine.AttnFuncWithCPAndKVP2P.forward")
return out_ret
@staticmethod
def backward(ctx, dout):
# pylint: disable=missing-function-docstring
nvtx_range_push("transformer_engine.AttnFuncWithCPAndKVP2P.backward")
cp_size_a2a = ctx.cp_size_a2a
rank_a2a = ctx.rank_a2a
......@@ -3602,6 +3609,7 @@ class AttnFuncWithCPAndKVP2P(torch.autograd.Function):
dq = ctx.dQKV_quantizer.create_tensor_from_data(dq, ctx.qkv_dtype)
dk = ctx.dQKV_quantizer.create_tensor_from_data(dk, ctx.qkv_dtype)
dv = ctx.dQKV_quantizer.create_tensor_from_data(dv, ctx.qkv_dtype)
nvtx_range_pop("transformer_engine.AttnFuncWithCPAndKVP2P.backward")
return (
None,
......@@ -3688,6 +3696,7 @@ class AttnFuncWithCPAndKVAllGather(torch.autograd.Function):
cp_stream,
):
# pylint: disable=missing-function-docstring
nvtx_range_push("transformer_engine.AttnFuncWithCPAndKVAllGather.forward")
if softmax_scale is None:
softmax_scale = q.shape[-1] ** (-0.5)
......@@ -3904,11 +3913,13 @@ class AttnFuncWithCPAndKVAllGather(torch.autograd.Function):
ctx.attn_mask_type = attn_mask_type
ctx.deterministic = deterministic
ctx.use_fused_attention = use_fused_attention
nvtx_range_pop("transformer_engine.AttnFuncWithCPAndKVAllGather.forward")
return out
@staticmethod
def backward(ctx, dout):
# pylint: disable=missing-function-docstring
nvtx_range_push("transformer_engine.AttnFuncWithCPAndKVAllGather.backward")
cp_size = get_distributed_world_size(ctx.cp_group)
rank = get_distributed_rank(ctx.cp_group)
......@@ -4092,6 +4103,7 @@ class AttnFuncWithCPAndKVAllGather(torch.autograd.Function):
dq = dq.view(*dq.shape[:seq_dim], -1, *dq.shape[(seq_dim + 2) :])
dk = dk.movedim(0, seq_dim).contiguous()
dv = dv.movedim(0, seq_dim).contiguous()
nvtx_range_pop("transformer_engine.AttnFuncWithCPAndKVAllGather.backward")
return (
None,
......@@ -4151,6 +4163,7 @@ class AttnFuncWithCPAndQKVOA2A(torch.autograd.Function):
quantizers,
):
# pylint: disable=missing-function-docstring
nvtx_range_push("transformer_engine.AttnFuncWithCPAndQKVOA2A.forward")
if softmax_scale is None:
softmax_scale = q.shape[-1] ** (-0.5)
......@@ -4403,11 +4416,13 @@ class AttnFuncWithCPAndQKVOA2A(torch.autograd.Function):
ctx.fp8_meta = fp8_meta
ctx.is_input_fp8 = is_input_fp8
ctx.is_output_fp8 = is_output_fp8
nvtx_range_pop("transformer_engine.AttnFuncWithCPAndQKVOA2A.forward")
return out_ret
@staticmethod
def backward(ctx, dout):
# pylint: disable=missing-function-docstring
nvtx_range_push("transformer_engine.AttnFuncWithCPAndQKVOA2A.backward")
cp_size = get_distributed_world_size(ctx.cp_group)
(
......@@ -4592,6 +4607,7 @@ class AttnFuncWithCPAndQKVOA2A(torch.autograd.Function):
dv = ctx.dQKV_quantizer.create_tensor_from_data(dv, fake_dtype=dout_dtype)
if not ctx.is_input_fp8:
dq, dk, dv = [x.dequantize() for x in [dq, dk, dv]]
nvtx_range_pop("transformer_engine.AttnFuncWithCPAndQKVOA2A.backward")
return (
None,
......
transformer_engine/pytorch/module/layernorm_linear.py
View file @ 49a4535d
......@@ -24,12 +24,14 @@ from .base import (
)
from ..fp8 import FP8GlobalStateManager
from ..utils import (
assert_dim_for_fp8_exec,
cast_if_needed,
clear_tensor_data,
divide,
get_default_init_method,
init_method_constant,
cast_if_needed,
assert_dim_for_fp8_exec,
clear_tensor_data,
nvtx_range_pop,
nvtx_range_push,
requires_grad,
)
from ..distributed import (
......@@ -112,6 +114,12 @@ class _LayerNormLinear(torch.autograd.Function):
skip_fp8_weight_update: bool,
) -> Union[Tuple[torch.Tensor, ...], torch.Tensor]:
# pylint: disable=missing-function-docstring
# NVTX label for profiling
nvtx_label = "transformer_engine._LayerNormLinear.forward"
if ub_name is not None:
nvtx_label = f"{nvtx_label}.{ub_name}"
# Make sure input dimensions are compatible
out_features, in_features = weight.shape
inp_shape = inp.shape
......@@ -121,10 +129,12 @@ class _LayerNormLinear(torch.autograd.Function):
assert_dim_for_fp8_exec(inputmat, weight)
# Cast for native AMP
nvtx_range_push(f"{nvtx_label}.norm_input_cast")
inputmat = cast_if_needed(inputmat, activation_dtype)
ln_weight = cast_if_needed(ln_weight, activation_dtype)
if ln_bias is not None:
ln_bias = cast_if_needed(ln_bias, activation_dtype)
nvtx_range_pop(f"{nvtx_label}.norm_input_cast")
tp_world_size = get_distributed_world_size(tp_group)
ub_overlap_ag_fprop = (
......@@ -175,6 +185,7 @@ class _LayerNormLinear(torch.autograd.Function):
)
# Apply normalization
nvtx_range_push(f"{nvtx_label}.norm")
ln_out, mu, rsigma = apply_normalization(
inputmat,
ln_out,
......@@ -188,9 +199,11 @@ class _LayerNormLinear(torch.autograd.Function):
zero_centered_gamma,
)
ln_out_return = ln_out if return_layernorm_output else None
nvtx_range_pop(f"{nvtx_label}.norm")
# Prepare GEMM input
# Note: Cast to expected dtype and perform tensor-parallel communication
nvtx_range_push(f"{nvtx_label}.gemm_input_cast_comm")
if with_input_all_gather and not ub_overlap_ag_fprop:
with_quantized_all_gather = fp8
if return_layernorm_output and return_layernorm_output_gathered:
......@@ -217,6 +230,7 @@ class _LayerNormLinear(torch.autograd.Function):
elif backward_needs_input:
ln_out.update_usage(rowwise_usage=True, columnwise_usage=True)
ln_out_total = ln_out
nvtx_range_pop(f"{nvtx_label}.gemm_input_cast_comm")
# Cast weight to expected dtype
weightmat = weight
......@@ -275,6 +289,7 @@ class _LayerNormLinear(torch.autograd.Function):
assert ub_obj.is_fp8_ubuf(), "AG overlap with FP8 GEMM inputs requires FP8 buffer."
ln_out_total = ub_obj.get_buffer(input_quantizer)
nvtx_range_push(f"{nvtx_label}.gemm")
out, *_, rs_out = general_gemm(
weightmat,
ln_out_total,
......@@ -287,6 +302,8 @@ class _LayerNormLinear(torch.autograd.Function):
ub_type=ub_type,
extra_output=rs_out,
)
nvtx_range_pop(f"{nvtx_label}.gemm")
if not weight.requires_grad:
if not return_layernorm_output:
ln_out = ln_out_total = None
......@@ -307,6 +324,7 @@ class _LayerNormLinear(torch.autograd.Function):
# Scatter intermediate/activation tensors saved for the backward pass
# NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already
# shards/unshards the base weights so we don't do it ourselves
nvtx_range_push(f"{nvtx_label}.fsdp_scatter")
ctx.fsdp_group = fsdp_group
ctx.fsdp_shapes = _fsdp_scatter_tensors(
fsdp_group,
......@@ -315,6 +333,7 @@ class _LayerNormLinear(torch.autograd.Function):
weightmat if quantized_weight else None,
ln_out if weight.requires_grad else None,
)
nvtx_range_pop(f"{nvtx_label}.fsdp_scatter")
tensors_to_save, tensor_objects = prepare_for_saving(
inputmat,
......@@ -372,10 +391,12 @@ class _LayerNormLinear(torch.autograd.Function):
if ub_overlap_rs_fprop:
out = rs_out
elif parallel_mode == "row":
nvtx_range_push(f"{nvtx_label}.row_parallel_comm")
if sequence_parallel:
out, _ = reduce_scatter_along_first_dim(out, tp_group)
elif tensor_parallel:
out, _ = allreduce(out, tp_group)
nvtx_range_pop(f"{nvtx_label}.row_parallel_comm")
# [*, in_features] -> [*, out_features] except first dimension changes for SP
out = out.view(-1, *inp_shape[1:-1], out_features)
......@@ -394,6 +415,11 @@ class _LayerNormLinear(torch.autograd.Function):
) -> Tuple[Union[torch.Tensor, None], ...]:
# pylint: disable=missing-function-docstring
# NVTX label for profiling
nvtx_label = "transformer_engine._LayerNormLinear.backward"
if ctx.ub_name is not None:
nvtx_label = f"{nvtx_label}.{ctx.ub_name}"
with torch.cuda.nvtx.range("_LayerNormLinear_backward"):
if (
ctx.fp8
......@@ -433,6 +459,7 @@ class _LayerNormLinear(torch.autograd.Function):
# Gather intermediate/activation tensors if needed
# NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already
# shards/unshards the base weights so we don't do it ourselves
nvtx_range_push(f"{nvtx_label}.fsdp_gather")
_fsdp_gather_tensors(
ctx.fsdp_group,
ctx.fsdp_shapes,
......@@ -441,6 +468,7 @@ class _LayerNormLinear(torch.autograd.Function):
weight if ctx.fp8 and ctx.quantized_weight else None,
ln_out,
)
nvtx_range_pop(f"{nvtx_label}.fsdp_gather")
# For CPU offloading, we offloaded weight and weight.main_grad to different tensors,
# we need to connect them into one.
......@@ -515,12 +543,14 @@ class _LayerNormLinear(torch.autograd.Function):
if ctx.fp8:
quantizer = ctx.input_quantizer
quantizer.set_usage(rowwise=True, columnwise=True)
nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input")
ln_out_total, ln_out_total_work = gather_along_first_dim(
ln_out,
ctx.tp_group,
async_op=True,
quantizer=quantizer,
)
nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input")
else:
ln_out_total = ln_out
......@@ -536,6 +566,7 @@ class _LayerNormLinear(torch.autograd.Function):
if ctx.grad_input_quantizer is not None:
ctx.grad_input_quantizer.set_usage(rowwise=True, columnwise=False)
nvtx_range_push(f"{nvtx_label}.dgrad_gemm")
dgrad, *_ = general_gemm(
weight,
grad_output,
......@@ -551,12 +582,14 @@ class _LayerNormLinear(torch.autograd.Function):
extra_output=rs_out,
bulk_overlap=ctx.ub_bulk_dgrad,
)
nvtx_range_pop(f"{nvtx_label}.dgrad_gemm")
# Launch tensor-parallel communication
dgrad_work = None
if ctx.ub_overlap_rs_dgrad:
dgrad = rs_out
elif ctx.parallel_mode == "column" and not ctx.ub_bulk_wgrad:
nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad")
if ctx.sequence_parallel:
if ctx.return_layernorm_output and ctx.return_layernorm_output_gathered:
dgrad = dgrad + grad_outputs[1].view_as(dgrad)
......@@ -567,6 +600,7 @@ class _LayerNormLinear(torch.autograd.Function):
)
else:
dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True)
nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad")
# Compute grad weight tensor
wgrad = None
......@@ -603,6 +637,7 @@ class _LayerNormLinear(torch.autograd.Function):
# wgrad GEMM
# Note: Fuse with bgrad computation if needed
nvtx_range_push(f"{nvtx_label}.wgrad_gemm")
wgrad, grad_bias_, *_, rs_out = general_gemm(
ln_out_total,
grad_output,
......@@ -621,6 +656,7 @@ class _LayerNormLinear(torch.autograd.Function):
extra_output=rs_out,
bulk_overlap=ctx.ub_bulk_wgrad,
)
nvtx_range_pop(f"{nvtx_label}.wgrad_gemm")
if ctx.ub_bulk_wgrad:
if ub_obj_wgrad.is_fp8_ubuf():
......@@ -657,6 +693,7 @@ class _LayerNormLinear(torch.autograd.Function):
# Norm gradient
dgamma = None
dbeta = None
nvtx_range_push(f"{nvtx_label}.norm")
if ctx.normalization == "LayerNorm":
dgrad, dgamma, dbeta = tex.layernorm_bwd(
dgrad,
......@@ -679,6 +716,7 @@ class _LayerNormLinear(torch.autograd.Function):
)
dgrad = dgrad.reshape(inputmat.size())
dbeta = None
nvtx_range_pop(f"{nvtx_label}.norm")
clear_tensor_data(mu)
clear_tensor_data(rsigma)
......
transformer_engine/pytorch/module/linear.py
View file @ 49a4535d
......@@ -22,12 +22,14 @@ from .base import (
from ._common import noop_cat, _fix_gathered_fp8_transpose
from ..fp8 import FP8GlobalStateManager
from ..utils import (
divide,
cast_if_needed,
clear_tensor_data,
divide,
init_method_constant,
requires_grad,
non_tn_fp8_gemm_supported,
nvtx_range_pop,
nvtx_range_push,
requires_grad,
)
from ..distributed import (
set_tensor_model_parallel_attributes,
......@@ -100,6 +102,11 @@ class _Linear(torch.autograd.Function):
) -> torch.Tensor:
# pylint: disable=missing-function-docstring
# NVTX label for profiling
nvtx_label = "transformer_engine._Linear.forward"
if ub_name is not None:
nvtx_label = f"{nvtx_label}.{ub_name}"
# Make sure input dimensions are compatible
out_features, in_features = weight.shape
inp_shape = inp.shape
......@@ -110,6 +117,7 @@ class _Linear(torch.autograd.Function):
# Prepare input tensor
# Note: Cast to expected dtype and perform tensor-parallel communication
nvtx_range_push(f"{nvtx_label}.input_cast_comm")
inputmat = inp
inputmat_total = None
with_input_all_gather_nccl = (
......@@ -153,6 +161,7 @@ class _Linear(torch.autograd.Function):
inputmat_total, _ = gather_along_first_dim(inputmat, tp_group)
else:
inputmat_total = inputmat
nvtx_range_pop(f"{nvtx_label}.input_cast_comm")
# Cast weight to expected dtype
weightmat = weight
......@@ -216,6 +225,7 @@ class _Linear(torch.autograd.Function):
ub_obj.copy_into_buffer(inputmat_total, input_quantizer, local_chunk=True)
inputmat_total = ub_obj.get_buffer(input_quantizer)
nvtx_range_push(f"{nvtx_label}.gemm")
out, *_, rs_out = general_gemm(
weightmat,
inputmat_total,
......@@ -228,6 +238,7 @@ class _Linear(torch.autograd.Function):
ub_type=ub_type,
extra_output=rs_out,
)
nvtx_range_pop(f"{nvtx_label}.gemm")
if is_grad_enabled:
saved_inputmat = None
......@@ -244,12 +255,14 @@ class _Linear(torch.autograd.Function):
# Scatter intermediate/activation tensors saved for the backward pass
# NOTE: FSDP sharding is not valid for models initialized with primary Fp8 weights
nvtx_range_push(f"{nvtx_label}.fsdp_scatter")
ctx.fsdp_group = fsdp_group
ctx.fsdp_shapes = _fsdp_scatter_tensors(
fsdp_group,
saved_inputmat,
weightmat if fp8 and not isinstance(weight, QuantizedTensor) else None,
)
nvtx_range_pop(f"{nvtx_label}.fsdp_scatter")
# TODO(ksivamani): Check memory usage
tensors_to_save, tensor_objects = prepare_for_saving(
......@@ -299,10 +312,12 @@ class _Linear(torch.autograd.Function):
if ub_overlap_rs_fprop:
out = rs_out
elif parallel_mode == "row":
nvtx_range_push(f"{nvtx_label}.row_parallel_comm")
if sequence_parallel:
out, _ = reduce_scatter_along_first_dim(out, tp_group)
elif tensor_parallel:
out, _ = allreduce(out, tp_group)
nvtx_range_pop(f"{nvtx_label}.row_parallel_comm")
out = out.view(-1, *inp_shape[1:-1], out_features)
return out
......@@ -311,6 +326,11 @@ class _Linear(torch.autograd.Function):
def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
# pylint: disable=missing-function-docstring
# NVTX label for profiling
nvtx_label = "transformer_engine._Linear.backward"
if ctx.ub_name is not None:
nvtx_label = f"{nvtx_label}.{ctx.ub_name}"
with torch.cuda.nvtx.range("_Linear_backward"):
if (
ctx.fp8
......@@ -347,12 +367,14 @@ class _Linear(torch.autograd.Function):
# Gather intermediate/activation tensors if needed
# NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already
# shards/unshards the base weights so we don't do it ourselves
nvtx_range_push(f"{nvtx_label}.fsdp_gather")
_fsdp_gather_tensors(
ctx.fsdp_group,
ctx.fsdp_shapes,
inputmat,
weight_fp8,
)
nvtx_range_pop(f"{nvtx_label}.fsdp_gather")
ctx.ub_obj_gradout = None
ub_obj_dgrad = None
......@@ -424,12 +446,14 @@ class _Linear(torch.autograd.Function):
if ctx.fp8:
quantizer = ctx.input_quantizer
quantizer.set_usage(rowwise=True, columnwise=True)
nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input")
inputmat_total, inputmat_total_work = gather_along_first_dim(
inputmat,
ctx.tp_group,
async_op=True,
quantizer=quantizer,
)
nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input")
else:
inputmat_total = inputmat
......@@ -451,6 +475,7 @@ class _Linear(torch.autograd.Function):
ctx.grad_input_quantizer.set_usage(rowwise=True, columnwise=False)
# dgrad GEMM
nvtx_range_push(f"{nvtx_label}.dgrad_gemm")
dgrad, *_, rs_out = general_gemm(
weight_fp8,
grad_output,
......@@ -466,11 +491,13 @@ class _Linear(torch.autograd.Function):
extra_output=rs_out,
bulk_overlap=ctx.ub_bulk_dgrad,
)
nvtx_range_pop(f"{nvtx_label}.dgrad_gemm")
# Launch tensor-parallel communication
if ctx.ub_overlap_rs_dgrad:
dgrad = rs_out
elif ctx.parallel_mode == "column" and not ctx.ub_bulk_wgrad:
nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad")
if ctx.sequence_parallel:
dgrad, dgrad_work = reduce_scatter_along_first_dim(
dgrad,
......@@ -479,6 +506,7 @@ class _Linear(torch.autograd.Function):
)
else:
dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True)
nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad")
# Compute grad weight tensor
wgrad = None
......@@ -515,6 +543,7 @@ class _Linear(torch.autograd.Function):
# wgrad GEMM
# Note: Fuse with bgrad computation if needed
nvtx_range_push(f"{nvtx_label}.wgrad_gemm")
wgrad, grad_bias_, _, rs_out = general_gemm(
inputmat_total,
grad_output,
......@@ -533,6 +562,7 @@ class _Linear(torch.autograd.Function):
extra_output=rs_out,
bulk_overlap=ctx.ub_bulk_wgrad,
)
nvtx_range_pop(f"{nvtx_label}.wgrad_gemm")
if ctx.ub_bulk_wgrad:
if ub_obj_wgrad.is_fp8_ubuf():
......
transformer_engine/pytorch/utils.py
View file @ 49a4535d
......@@ -6,6 +6,7 @@
from __future__ import annotations
import functools
import math
import os
from typing import Any, Callable, List, Optional, Tuple
import torch
......@@ -326,3 +327,62 @@ def round_up_to_nearest_multiple(value, multiple):
if multiple == 0:
raise ValueError("multiple cannot be zero.")
return ((value + multiple - 1) // multiple) * multiple
@functools.lru_cache(maxsize=None)
def _nvtx_enabled() -> bool:
"""Check if NVTX range profiling is enabled"""
return bool(int(os.getenv("NVTE_NVTX_ENABLED", "0")))
# Messages associated with active NVTX ranges
_nvtx_range_messages: list[str] = []
def nvtx_range_push(msg: str) -> None:
"""Push NVTX range onto stack, if NVTX range profiling is enabled
Set `NVTE_NVTX_ENABLED=1` in the environment to enable NVTX range
profiling.
Parameters
----------
msg: str
Message to associate with range
"""
if not _nvtx_enabled():
return
_nvtx_range_messages.append(msg)
torch.cuda.nvtx.range_push(msg)
def nvtx_range_pop(msg: Optional[str] = None) -> None:
"""Pop NVTX range from stack, if NVTX range profiling is enabled
Set `NVTE_NVTX_ENABLED=1` in the environment to enable NVTX range
profiling.
Parameters
----------
msg: str, optional
Message associated with range
"""
# Return immediately if NVTX range profiling is not enabled
if not _nvtx_enabled():
return
# Update list of NVTX range messages and check for consistency
if not _nvtx_range_messages:
raise RuntimeError("Attempted to pop NVTX range from empty stack")
last_msg = _nvtx_range_messages.pop()
if msg is not None and msg != last_msg:
raise ValueError(
f"Attempted to pop NVTX range from stack with msg={msg}, "
f"but last range has msg={last_msg}"
)
# Pop NVTX range
torch.cuda.nvtx.range_pop()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment