diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 2a486450a20..524fa639fac 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -35,6 +35,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_recipe.xml $TE_P python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_custom_recipe.xml $TE_PATH/tests/pytorch/test_custom_recipe.py || test_fail "test_custom_recipe.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_deferred_init.xml $TE_PATH/tests/pytorch/test_deferred_init.py || test_fail "test_deferred_init.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/test_numerics.py || test_fail "test_numerics.py" +PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_tied_delayed_wgrad.xml $TE_PATH/tests/pytorch/test_tied_delayed_wgrad.py || test_fail "test_tied_delayed_wgrad.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cuda_graphs.xml $TE_PATH/tests/pytorch/test_cuda_graphs.py || test_fail "test_cuda_graphs.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_jit.xml $TE_PATH/tests/pytorch/test_jit.py || test_fail "test_jit.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_rope.xml $TE_PATH/tests/pytorch/test_fused_rope.py || test_fail "test_fused_rope.py" diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 8f85f57f32b..a26e5d9a848 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -753,6 +753,180 @@ def test_make_graphed_callables_with_kwargs( assert_all_equal(outputs, graph_outputs) +@pytest.mark.parametrize("with_order", (False, True)) +@pytest.mark.parametrize("set_to_none", (False, True)) +@pytest.mark.parametrize("bias", (False, True)) +@pytest.mark.parametrize("dtype", (torch.float32, torch.bfloat16)) +def test_make_graphed_callables_with_delayed_wgrad(with_order, set_to_none, bias, dtype): + """Delayed replay survives zero_grad and matches independent optimizer steps.""" + reset_rng_states() + hidden_size, batch_size = 16, 8 + module = Linear( + hidden_size, + hidden_size, + bias=bias, + params_dtype=dtype, + delay_wgrad_compute=True, + fuse_wgrad_accumulation=False, + ) + reference = torch.nn.Linear(hidden_size, hidden_size, bias=bias, device="cuda", dtype=dtype) + with torch.no_grad(): + module.weight.fill_(0.25) + reference.weight.copy_(module.weight) + if bias: + module.bias.fill_(0.125) + reference.bias.copy_(module.bias) + graph_options = {} + if with_order: + graph_options = {"_order": [1, -1, -1.5], "_num_layers_per_chunk": [1]} + graphed = make_graphed_callables( + module, + (torch.ones(batch_size, hidden_size, device="cuda", dtype=dtype, requires_grad=True),), + **graph_options, + ) + optimizers = [torch.optim.SGD(model.parameters(), lr=1 / 64) for model in (module, reference)] + saved_grads = [] + try: + for _ in range(3): + for optimizer in optimizers: + optimizer.zero_grad(set_to_none=set_to_none) + # Bound the ordinary-allocator reuse pressure and check every buffer. + guards = [torch.full_like(module.weight, 17 + index) for index in range(32)] + input_ = torch.randint(-2, 3, (batch_size, hidden_size), device="cuda").to(dtype) / 4 + graph_input = input_.detach().requires_grad_(True) + reference_input = input_.detach().clone().requires_grad_(True) + grad_output = ( + torch.randint(-2, 3, (batch_size, hidden_size), device="cuda").to(dtype) / 8 + ) + output, expected = graphed(graph_input), reference(reference_input) + torch.testing.assert_close(output, expected, rtol=0, atol=0) + output.backward(grad_output) + expected.backward(grad_output) + graphed.backward_dw() + for index, guard in enumerate(guards): + torch.testing.assert_close( + guard, torch.full_like(guard, 17 + index), rtol=0, atol=0 + ) + torch.testing.assert_close(graph_input.grad, reference_input.grad, rtol=0, atol=0) + for param, reference_param in zip(module.parameters(), reference.parameters()): + assert param.grad is not None + torch.testing.assert_close(param.grad, reference_param.grad, rtol=0, atol=0) + if set_to_none: + for grad, snapshot in saved_grads: + torch.testing.assert_close(grad, snapshot, rtol=0, atol=0) + saved_grads.extend( + (param.grad, param.grad.clone()) for param in module.parameters() + ) + for optimizer in optimizers: + optimizer.step() + for param, reference_param in zip(module.parameters(), reference.parameters()): + torch.testing.assert_close(param, reference_param, rtol=0, atol=0) + finally: + graphed.reset() + + +@pytest.mark.parametrize("with_order", (False, True)) +@pytest.mark.parametrize("set_to_none", (False, True)) +@pytest.mark.parametrize("tied", (False, True)) +def test_graphed_delayed_wgrad_accumulates_microbatches(with_order, set_to_none, tied): + """Separate delayed graphs contribute to shared weights across microbatches.""" + reset_rng_states() + modules = tuple( + Linear(16, 16, bias=False, params_dtype=torch.float32, delay_wgrad_compute=True) + for _ in range(2) + ) + reference = torch.nn.Sequential( + torch.nn.Linear(16, 16, bias=False, device="cuda"), + torch.nn.Linear(16, 16, bias=False, device="cuda"), + ) + if tied: + modules[1].weight = modules[0].weight + reference[1].weight = reference[0].weight + with torch.no_grad(): + for module, reference_module in zip(modules, reference): + module.weight.fill_(0.25) + reference_module.weight.copy_(module.weight) + graph_options = {} + if with_order: + graph_options = {"_order": [1, 2, -2, -2.5, -1, -1.5], "_num_layers_per_chunk": [1, 1]} + graphed = make_graphed_callables( + modules, + tuple((torch.ones(8, 16, device="cuda", requires_grad=True),) for _ in modules), + **graph_options, + ) + optimizer = torch.optim.SGD(reference.parameters(), lr=1 / 64) + parameters = tuple(dict.fromkeys(param for module in modules for param in module.parameters())) + graph_optimizer = torch.optim.SGD(parameters, lr=1 / 64) + try: + for _ in range(3): + optimizer.zero_grad(set_to_none=set_to_none) + graph_optimizer.zero_grad(set_to_none=set_to_none) + for _ in range(2): + input_ = torch.randint(-2, 3, (8, 16), device="cuda").float() / 4 + graph_input = input_.detach().requires_grad_(True) + reference_input = input_.detach().clone().requires_grad_(True) + output = graphed[1](graphed[0](graph_input)) + expected = reference(reference_input) + torch.testing.assert_close(output, expected, rtol=0, atol=0) + output.sum().backward() + expected.sum().backward() + for callable_ in reversed(graphed): + callable_.backward_dw() + torch.testing.assert_close(graph_input.grad, reference_input.grad, rtol=0, atol=0) + for module, reference_module in zip(modules, reference): + assert module.weight.grad is not None + torch.testing.assert_close( + module.weight.grad, reference_module.weight.grad, rtol=0, atol=0 + ) + graph_optimizer.step() + optimizer.step() + for module, reference_module in zip(modules, reference): + torch.testing.assert_close(module.weight, reference_module.weight, rtol=0, atol=0) + finally: + reset_graphs(graphed) + + +@pytest.mark.parametrize("with_order", (False, True)) +def test_graphed_delayed_wgrad_preserves_main_grad(with_order): + """Fused accumulation continues to write main_grad across optimizer boundaries.""" + reset_rng_states() + module = Linear( + 16, + 16, + bias=False, + params_dtype=torch.float32, + delay_wgrad_compute=True, + fuse_wgrad_accumulation=True, + ) + module.weight.main_grad = torch.zeros_like(module.weight) + reference = torch.nn.Linear(16, 16, bias=False, device="cuda") + with torch.no_grad(): + module.weight.fill_(0.25) + reference.weight.copy_(module.weight) + graph_options = {} + if with_order: + graph_options = {"_order": [1, -1, -1.5], "_num_layers_per_chunk": [1]} + graphed = make_graphed_callables( + module, (torch.ones(8, 16, device="cuda", requires_grad=True),), **graph_options + ) + try: + for _ in range(2): + module.zero_grad() + module.weight.main_grad.zero_() + reference.zero_grad() + for _ in range(2): + input_ = torch.randint(-2, 3, (8, 16), device="cuda").float() / 4 + graphed(input_.detach().requires_grad_(True)).sum().backward() + reference(input_.detach().clone().requires_grad_(True)).sum().backward() + graphed.backward_dw() + assert module.weight.grad is None + torch.testing.assert_close( + module.weight.main_grad, reference.weight.grad, rtol=0, atol=0 + ) + finally: + graphed.reset() + + def test_make_graphed_callables_returns_owned_parameter_grads() -> None: """Parameter grads returned from graph replay must not alias static graph buffers.""" reset_rng_states() diff --git a/tests/pytorch/test_tied_delayed_wgrad.py b/tests/pytorch/test_tied_delayed_wgrad.py new file mode 100644 index 00000000000..75a53a9b695 --- /dev/null +++ b/tests/pytorch/test_tied_delayed_wgrad.py @@ -0,0 +1,334 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Delayed weight gradients accumulate for tied parameters and microbatches.""" + +import pytest +import torch +from torch import nn + +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import Float8BlockScaling +from transformer_engine.pytorch.module.grouped_linear import is_module_grouped_tensor_path_supported +from transformer_engine.pytorch.ops.basic.grouped_linear import ( + is_op_fuser_grouped_tensor_path_supported, +) +from transformer_engine.pytorch.quantization import is_fp8_block_scaling_available + + +def _build_two_linears(delay, *, bias=False, device="cuda"): + """Build two BF16 Linear layers with independent parameters.""" + hidden_size = 16 + return nn.Sequential( + te.Linear( + hidden_size, + hidden_size, + bias=bias, + params_dtype=torch.bfloat16, + device=device, + delay_wgrad_compute=delay, + fuse_wgrad_accumulation=False, + ), + te.Linear( + hidden_size, + hidden_size, + bias=bias, + params_dtype=torch.bfloat16, + device=device, + delay_wgrad_compute=delay, + fuse_wgrad_accumulation=False, + ), + ) + + +@pytest.mark.parametrize("bias", [False, True]) +@pytest.mark.parametrize("set_to_none", [False, True]) +@pytest.mark.parametrize("microbatches", [1, 2]) +@pytest.mark.parametrize("backward_dw_order", [(0, 1), (1, 0)]) +def test_tied_linear_gradients(bias, set_to_none, microbatches, backward_dw_order): + """Tied weight and bias gradients match regular backward after either reset.""" + torch.manual_seed(9) + delayed = _build_two_linears(True, bias=bias) + regular = _build_two_linears(False, bias=bias) + regular.load_state_dict(delayed.state_dict()) + for model in (delayed, regular): + model[1].weight = model[0].weight + if bias: + model[1].bias = model[0].bias + hidden_size, batch_size = 16, 2 + optimizers = [torch.optim.SGD(model.parameters(), lr=0.0) for model in (delayed, regular)] + for _ in range(2): + for optimizer in optimizers: + optimizer.zero_grad(set_to_none=set_to_none) + x = torch.randn(batch_size * microbatches, hidden_size, device="cuda", dtype=torch.bfloat16) + for model in (delayed, regular): + for microbatch in x.chunk(microbatches): + model(microbatch.detach().clone().requires_grad_(True)).float().sum().backward() + for _ in range(microbatches): + for index in backward_dw_order: + delayed[index].backward_dw() + # Microbatch contributions can accumulate in a different BF16 rounding order. + tolerance = 1e-2 if microbatches > 1 else 0.0 + torch.testing.assert_close( + delayed[0].weight.grad, regular[0].weight.grad, rtol=tolerance, atol=tolerance + ) + assert torch.count_nonzero(delayed[0].weight.grad) > 0 + if bias: + torch.testing.assert_close( + delayed[0].bias.grad, regular[0].bias.grad, rtol=tolerance, atol=tolerance + ) + for optimizer in optimizers: + optimizer.step() + + +def test_non_tied_linear_gradients(): + """Independent delayed parameters match regular backward.""" + torch.manual_seed(9) + delayed, regular = _build_two_linears(True), _build_two_linears(False) + regular.load_state_dict(delayed.state_dict()) + x = torch.randn(2, 16, device="cuda", dtype=torch.bfloat16) + for model in (delayed, regular): + model(x.detach().clone().requires_grad_(True)).float().sum().backward() + for index in (0, 1): + delayed[index].backward_dw() + torch.testing.assert_close( + delayed[index].weight.grad, regular[index].weight.grad, rtol=0, atol=0 + ) + + +def test_tied_linear_matches_pytorch(): + """A plain PyTorch Linear provides an independent tied-weight reference.""" + torch.manual_seed(9) + model = _build_two_linears(True) + model[1].weight = model[0].weight + reference = nn.Linear(16, 16, bias=False, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + reference.weight.copy_(model[0].weight) + x = torch.randn(2, 16, device="cuda", dtype=torch.bfloat16) + model(x.detach().clone().requires_grad_(True)).float().sum().backward() + for layer in model: + layer.backward_dw() + reference_x = x.detach().clone().requires_grad_(True) + reference(reference(reference_x)).float().sum().backward() + torch.testing.assert_close(model[0].weight.grad, reference.weight.grad, rtol=0, atol=0) + + +@pytest.mark.parametrize("set_to_none", [False, True]) +@pytest.mark.parametrize("microbatches", [1, 2]) +@pytest.mark.parametrize("backward_dw_order", [(0, 1), (1, 0)]) +@pytest.mark.parametrize("fp8_block_scaling", [False, True]) +def test_tied_layernorm_mlp_gradients( + set_to_none, microbatches, backward_dw_order, fp8_block_scaling +): + """Both FCs accumulate delayed gradients without duplicating unfused bias grads.""" + if fp8_block_scaling: + supported, reason = is_fp8_block_scaling_available(return_reason=True) + if not supported: + pytest.skip(reason) + torch.manual_seed(9) + hidden_size, ffn_hidden_size, batch_size = 128, 256, 128 + delayed = nn.Sequential( + te.LayerNormMLP( + hidden_size, + ffn_hidden_size, + params_dtype=torch.bfloat16, + delay_wgrad_compute=True, + fuse_wgrad_accumulation=False, + ), + te.LayerNormMLP( + hidden_size, + ffn_hidden_size, + params_dtype=torch.bfloat16, + delay_wgrad_compute=True, + fuse_wgrad_accumulation=False, + ), + ) + regular = nn.Sequential( + te.LayerNormMLP( + hidden_size, + ffn_hidden_size, + params_dtype=torch.bfloat16, + delay_wgrad_compute=False, + fuse_wgrad_accumulation=False, + ), + te.LayerNormMLP( + hidden_size, + ffn_hidden_size, + params_dtype=torch.bfloat16, + delay_wgrad_compute=False, + fuse_wgrad_accumulation=False, + ), + ) + regular.load_state_dict(delayed.state_dict()) + names = ("fc1_weight", "fc2_weight", "fc1_bias", "fc2_bias") + for model in (delayed, regular): + for name in names: + setattr(model[1], name, getattr(model[0], name)) + optimizers = [torch.optim.SGD(model.parameters(), lr=0.0) for model in (delayed, regular)] + for _ in range(2): + for optimizer in optimizers: + optimizer.zero_grad(set_to_none=set_to_none) + x = torch.randn(batch_size * microbatches, hidden_size, device="cuda", dtype=torch.bfloat16) + for model in (delayed, regular): + for microbatch in x.chunk(microbatches): + with te.autocast(enabled=fp8_block_scaling, recipe=Float8BlockScaling()): + output = model(microbatch.detach().clone().requires_grad_(True)) + output.float().sum().backward() + for _ in range(microbatches): + for index in backward_dw_order: + delayed[index].backward_dw() + for name in names: + torch.testing.assert_close( + getattr(delayed[0], name).grad, + getattr(regular[0], name).grad, + rtol=1e-2, + atol=1e-2, + ) + assert torch.count_nonzero(delayed[0].fc1_weight.grad) > 0 + assert torch.count_nonzero(delayed[0].fc2_weight.grad) > 0 + for optimizer in optimizers: + optimizer.step() + + +@pytest.mark.parametrize("set_to_none", [False, True]) +@pytest.mark.parametrize("backward_dw_order", [(0, 1), (1, 0)]) +@pytest.mark.parametrize("packed", [False, True]) +def test_tied_grouped_linear_gradients(set_to_none, backward_dw_order, packed, monkeypatch): + """Module grouped weights and deferred discrete biases accumulate across two microbatches.""" + dtype = torch.bfloat16 + if packed and not is_module_grouped_tensor_path_supported(None, dtype): + pytest.skip("packed grouped weights require a supported GPU and cuBLASLt") + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1" if packed else "0") + torch.manual_seed(9) + num_groups, hidden_size, tokens_per_group, microbatches = 2, 16, 16, 2 + delayed = nn.ModuleList( + [ + te.GroupedLinear( + num_groups, + hidden_size, + hidden_size, + bias=not packed, + params_dtype=dtype, + fuse_wgrad_accumulation=False, + use_grouped_tensor=packed, + delay_wgrad_compute=True, + single_grouped_weight=packed, + ) + for _ in range(2) + ] + ) + regular = nn.ModuleList( + [ + te.GroupedLinear( + num_groups, + hidden_size, + hidden_size, + bias=not packed, + params_dtype=dtype, + fuse_wgrad_accumulation=False, + use_grouped_tensor=packed, + delay_wgrad_compute=False, + single_grouped_weight=packed, + ) + for _ in range(2) + ] + ) + regular.load_state_dict(delayed.state_dict()) + names = ("weight",) if packed else ("weight0", "weight1", "bias0", "bias1") + for model in (delayed, regular): + for name in names: + setattr(model[1], name, getattr(model[0], name)) + splits = [tokens_per_group] * num_groups + if packed: + splits = torch.tensor(splits, device="cuda", dtype=torch.int64) + for _ in range(2): + for model in (delayed, regular): + model.zero_grad(set_to_none=set_to_none) + x = torch.randn( + microbatches * tokens_per_group * num_groups, hidden_size, device="cuda", dtype=dtype + ) + for model in (delayed, regular): + for microbatch in x.chunk(microbatches): + output = model[0](microbatch.detach().clone().requires_grad_(True), splits) + model[1](output, splits).float().sum().backward() + for _ in range(microbatches): + for index in backward_dw_order: + delayed[index].backward_dw() + for name in names: + torch.testing.assert_close( + getattr(delayed[0], name).grad, + getattr(regular[0], name).grad, + rtol=1e-2, + atol=1e-2, + ) + assert torch.count_nonzero(getattr(delayed[0], names[0]).grad) > 0 + + +@pytest.mark.parametrize("set_to_none", [False, True]) +@pytest.mark.parametrize("backward_dw_order", [(0, 1), (1, 0)]) +@pytest.mark.parametrize("packed", [False, True]) +def test_tied_grouped_linear_ops_gradients(set_to_none, backward_dw_order, packed, monkeypatch): + """Fusible grouped operations accumulate delayed packed and discrete weight grads.""" + dtype = torch.bfloat16 + if packed and not is_op_fuser_grouped_tensor_path_supported(None, dtype): + pytest.skip("packed grouped weights require a supported GPU and cuBLASLt") + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1" if packed else "0") + torch.manual_seed(9) + num_groups, hidden_size, tokens_per_group, microbatches = 2, 16, 16, 2 + delayed = nn.ModuleList( + [ + te.ops.GroupedLinear( + num_groups, + hidden_size, + hidden_size, + bias=False, + dtype=dtype, + accumulate_into_main_grad=False, + delay_wgrad_compute=True, + single_grouped_weight=packed, + ) + for _ in range(2) + ] + ) + regular = nn.ModuleList( + [ + te.ops.GroupedLinear( + num_groups, + hidden_size, + hidden_size, + bias=False, + dtype=dtype, + accumulate_into_main_grad=False, + delay_wgrad_compute=False, + single_grouped_weight=packed, + ) + for _ in range(2) + ] + ) + regular.load_state_dict(delayed.state_dict()) + names = ("weight",) if packed else ("weight0", "weight1") + for model in (delayed, regular): + for name in names: + setattr(model[1], name, getattr(model[0], name)) + splits = torch.tensor([tokens_per_group] * num_groups, device="cuda", dtype=torch.int64) + for _ in range(2): + for model in (delayed, regular): + model.zero_grad(set_to_none=set_to_none) + x = torch.randn( + microbatches * tokens_per_group * num_groups, hidden_size, device="cuda", dtype=dtype + ) + for model in (delayed, regular): + for microbatch in x.chunk(microbatches): + output = model[0](microbatch.detach().clone().requires_grad_(True), splits) + model[1](output, splits).float().sum().backward() + for _ in range(microbatches): + for index in backward_dw_order: + delayed[index].backward_dw() + for name in names: + torch.testing.assert_close( + getattr(delayed[0], name).grad, + getattr(regular[0], name).grad, + rtol=1e-2, + atol=1e-2, + ) + assert torch.count_nonzero(getattr(delayed[0], names[0]).grad) > 0 diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index bc89a815af1..be88f730a18 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -529,6 +529,22 @@ def _returned_param_grad_clone_slots(static_grad_inputs, module_params): # Filter the TE modules that cudagraph can access. visited_te_modules = {} need_bwd_dw_graph = {} + per_callable_static_wgrads = [()] * len(flatten_sample_args) + + def _capture_backward_dw(func_idx, bwd_dw_graph): + """Capture deferred contributions without retaining ordinary-allocator grad targets.""" + modules = tuple( + module + for module in visited_te_modules[func_idx] + if hasattr(module, "need_backward_dw") and module.need_backward_dw() + ) + params = tuple(dict.fromkeys(param for module in modules for param in module.parameters())) + # Warmup grads belong to the ordinary allocator. Capture fresh results + # in the graph pool so zero_grad() cannot free an accumulation target. + with _none_grad_context_wrapper(params), _graph_context_wrapper(bwd_dw_graph, pool=mempool): + for module in modules: + module.backward_dw() + return tuple((param, param.grad) for param in params if param.grad is not None) def _run_warmup_forward(func_idx, func, callable_idx): """Run forward for one callable during warmup; returns flattened outputs.""" @@ -806,13 +822,9 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): "No module needs wgrad computation but get float in order" ) bwd_dw_graph = bwd_dw_graphs[per_callable_bwd_idx] - with _graph_context_wrapper(bwd_dw_graph, pool=mempool): - for module in visited_te_modules[per_callable_bwd_idx]: - if ( - hasattr(module, "need_backward_dw") - and module.need_backward_dw() - ): - module.backward_dw() + per_callable_static_wgrads[per_callable_bwd_idx] = _capture_backward_dw( + per_callable_bwd_idx, bwd_dw_graph + ) continue static_input_surface = per_callable_static_input_surfaces[per_callable_bwd_idx] @@ -1005,10 +1017,9 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): ) if need_bwd_dw_graph[bwd_idx]: - with _graph_context_wrapper(bwd_dw_graph, pool=mempool): - for module in visited_te_modules[bwd_idx]: - if hasattr(module, "need_backward_dw") and module.need_backward_dw(): - module.backward_dw() + per_callable_static_wgrads[bwd_idx] = _capture_backward_dw( + bwd_idx, bwd_dw_graph + ) # Constructs a tuple suitable for returning from Graphed.backward: # Pads out the actually-needed grads with Nones in gradient slots for inputs that # don't require grad. I couldn't think of a slick one-liner for this pattern. @@ -1237,6 +1248,7 @@ def make_graphed_attribute_functions(graph_idx, helpers): bwd_dw_graph = bwd_dw_graphs[graph_idx] need_bwd_dw = need_bwd_dw_graph.get(graph_idx, False) te_modules = visited_te_modules.get(graph_idx, set()) + static_wgrads = per_callable_static_wgrads[graph_idx] # Attach backward_dw as an attribute to the graphed callable. def backward_dw(): @@ -1244,6 +1256,15 @@ def backward_dw(): if need_bwd_dw: bwd_dw_graph.replay() + # Python assignments made during capture are not replayed. Give + # callers owned grads and accumulate tied/microbatch contributions + # before running their delayed-grad hooks. + for param, grad in static_wgrads: + if param.grad is None: + param.grad = grad.detach().clone() + else: + param.grad.add_(grad) + # Trigger the grad accumulation hook for wgrad graphs. for module in te_modules: if ( @@ -1254,7 +1275,7 @@ def backward_dw(): # Attach reset as an attribute to the graphed callable. def reset(): - nonlocal fwd_graph, bwd_graph, bwd_dw_graph, te_modules + nonlocal fwd_graph, bwd_graph, bwd_dw_graph, te_modules, static_wgrads for graph in (fwd_graph, bwd_graph, bwd_dw_graph): if graph is not None: @@ -1264,6 +1285,7 @@ def reset(): bwd_graph = None bwd_dw_graph = None te_modules = () + static_wgrads = () helpers.release_static_state() return backward_dw, reset diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 9ead1e73651..f7d9accb5f4 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -2005,11 +2005,19 @@ def backward_dw(self): (wgrad, bgrad), _ = self.wgrad_store.pop() if not self.fuse_wgrad_accumulation: weight_tensor = noop_cat(self._get_weight_tensors()) - weight_tensor.grad = wgrad.to(weight_tensor.dtype) + wgrad = wgrad.to(weight_tensor.dtype) + # Tied parameters and microbatches contribute to the same gradient. + if weight_tensor.grad is None: + weight_tensor.grad = wgrad + else: + weight_tensor.grad.add_(wgrad) if self.use_bias and bgrad is not None and bgrad.numel() != 0: bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names]) + bgrad = bgrad.to(bias_tensor.dtype) if bias_tensor.grad is None: - bias_tensor.grad = bgrad.to(bias_tensor.dtype) + bias_tensor.grad = bgrad + else: + bias_tensor.grad.add_(bgrad) del wgrad del bgrad self._trigger_wgrad_accumulation_and_reduce_hooks() diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 5c0dbf3ab28..d8778e7d8e1 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -2430,12 +2430,20 @@ def backward_dw(self): weight_params = self._get_weight_tensors() if not self.fuse_wgrad_accumulation: if self.single_grouped_weight: - weight_params[0].grad = wgrad_output.rowwise_data.view( + wgrad = wgrad_output.rowwise_data.view( self.num_gemms, self.out_features, self.in_features ).to(weight_params[0].dtype) + if weight_params[0].grad is None: + weight_params[0].grad = wgrad + else: + weight_params[0].grad.add_(wgrad) else: for i in range(self.num_gemms): - weight_params[i].grad = wgrad_output[i].to(weight_params[i].dtype) + wgrad = wgrad_output[i].to(weight_params[i].dtype) + if weight_params[i].grad is None: + weight_params[i].grad = wgrad + else: + weight_params[i].grad.add_(wgrad) has_grad_biases = [ grad_bias is not None and grad_bias.numel() != 0 for grad_bias in grad_biases_ ] @@ -2450,8 +2458,12 @@ def backward_dw(self): ) bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] for i in range(self.num_gemms): - if has_grad_biases[i] and bias_params[i].grad is None: - bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype) + if has_grad_biases[i]: + bgrad = grad_biases_[i].to(bias_params[i].dtype) + if bias_params[i].grad is None: + bias_params[i].grad = bgrad + else: + bias_params[i].grad.add_(bgrad) del grad_biases_ del wgrad_output del tensor_list diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index ad5864d4a60..a3d800667cf 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -3140,36 +3140,31 @@ def _get_weight_quantizers(self) -> List[Quantizer]: return [fc1_weight_quantizer, fc2_weight_quantizer] def backward_dw(self): - """ - Execute the delayed weight gradient computation. - This method is called after the main backward pass to compute weight gradients. - """ + """Execute delayed weight-gradient GEMMs and accumulate their contributions.""" if not self.need_backward_dw(): return with get_nvtx_range_context("_LayerNormMLP_wgrad"): - (fc2_wgrad, fc2_bias_grad_, *_), tensor_list_fc2 = self.wgrad_store.pop() - if self.use_bias and self.fc1_bias.grad is None: - (fc1_wgrad, fc1_bias_grad, *_), _ = self.wgrad_store.pop() - else: - (fc1_wgrad, *_), _ = self.wgrad_store.pop() - fc1_bias_grad = None + (fc2_wgrad, fc2_bias_grad_, *_), _ = self.wgrad_store.pop() + (fc1_wgrad, fc1_bias_grad, *_), _ = self.wgrad_store.pop() if self.use_bias: - if self.fc2_bias.grad is None: - if ( - self.fp8 - and FP8GlobalStateManager.get_fp8_recipe().float8_block_scaling() - and self.apply_bias - and not self.gemm_bias_unfused_add - ): - act_out = tensor_list_fc2[0] - # BGRAD not fused with GEMM for float8 blockwise gemm. - fc2_bias_grad_ = sum_bias_grad(act_out) - self.fc2_bias.grad = fc2_bias_grad_.to(self.fc2_bias.dtype) - if self.fc1_bias.grad is None: - self.fc1_bias.grad = fc1_bias_grad.to(self.fc1_bias.dtype) + # Unfused bias gradients are already returned by the main backward. + for bias, bgrad in ( + (self.fc2_bias, fc2_bias_grad_), + (self.fc1_bias, fc1_bias_grad), + ): + if bgrad is not None and bgrad.numel() != 0: + bgrad = bgrad.to(bias.dtype) + if bias.grad is None: + bias.grad = bgrad + else: + bias.grad.add_(bgrad) if not self.fuse_wgrad_accumulation: - self.fc2_weight.grad = fc2_wgrad.to(self.fc2_weight.dtype) - self.fc1_weight.grad = fc1_wgrad.to(self.fc1_weight.dtype) + for weight, wgrad in ((self.fc2_weight, fc2_wgrad), (self.fc1_weight, fc1_wgrad)): + wgrad = wgrad.to(weight.dtype) + if weight.grad is None: + weight.grad = wgrad + else: + weight.grad.add_(wgrad) del fc2_bias_grad_ del fc2_wgrad del fc1_wgrad diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 508895505fd..32114f55a61 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -364,17 +364,25 @@ def backward_dw(self) -> None: return if self.single_grouped_weight: if isinstance(grad_weights, list): - self.weight.grad = torch.stack(grad_weights, dim=0).to(self.weight.dtype) + wgrad = torch.stack(grad_weights, dim=0).to(self.weight.dtype) else: - self.weight.grad = grad_weights.rowwise_data.view( + wgrad = grad_weights.rowwise_data.view( self.num_groups, self.out_features, self.in_features, ).to(self.weight.dtype) + if self.weight.grad is None: + self.weight.grad = wgrad + else: + self.weight.grad.add_(wgrad) else: for group_idx in range(self.num_groups): w = getattr(self, f"weight{group_idx}") - w.grad = grad_weights[group_idx].to(w.dtype) + wgrad = grad_weights[group_idx].to(w.dtype) + if w.grad is None: + w.grad = wgrad + else: + w.grad.add_(wgrad) self._trigger_wgrad_accumulation_and_reduce_hooks() def _get_discrete_bias_tensors(self, dtype: torch.dtype) -> list[torch.Tensor]: