Tied weights do not accumulate with delayed wgrad

[wujingyue]

Describe the bug

Transformer Engine does not accumulate the gradient of a tied weight when
delay_wgrad_compute=True and fuse_wgrad_accumulation=False.

Two tied delayed te.Linear modules produce a different shared-weight gradient
than two equivalent tied normal te.Linear modules. The equivalent non-tied
control case passes.

Steps/Code to reproduce bug

import torch
import transformer_engine.pytorch as te
from torch import nn


def make_model(delay_wgrad_compute: bool) -> nn.Sequential:
    model = nn.Sequential(
        te.Linear(
            16, 16, bias=False, params_dtype=torch.bfloat16, device="cuda",
            delay_wgrad_compute=delay_wgrad_compute,
            fuse_wgrad_accumulation=False,
        ),
        te.Linear(
            16, 16, bias=False, params_dtype=torch.bfloat16, device="cuda",
            delay_wgrad_compute=delay_wgrad_compute,
            fuse_wgrad_accumulation=False,
        ),
    )
    model[1].weight = model[0].weight
    return model


torch.manual_seed(9)
delayed = make_model(True)
regular = make_model(False)
regular.load_state_dict(delayed.state_dict())

x = torch.randn(2, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True)

delayed(x).float().sum().backward()
regular(x).float().sum().backward()

delayed[0].backward_dw()
delayed[1].backward_dw()

torch.testing.assert_close(delayed[0].weight.grad, regular[0].weight.grad)

Expected behavior

The shared tied-weight gradient from the delayed model should match the shared
tied-weight gradient from the equivalent normal model.

Actual behavior

AssertionError: Tensor-likes are not close!

Mismatched elements: 255 / 256 (99.6%)
Greatest absolute difference: 0.396484375 at index (3, 7) (up to 1e-05 allowed)
Greatest relative difference: 22.875 at index (5, 12) (up to 0.016 allowed)

Environment details

• OS: Linux 6.8.0
• Python: 3.12.3
• PyTorch: 2.13.0a0+8145d630e8.nv26.06
• Transformer Engine: 2.17.1+4329ff84
• CUDA: 13.3
• GPU: NVIDIA RTX A6000
• Driver: 570.211.01

Additional context

This was found while adding MFSDP support for Transformer Engine delayed weight
gradients: NVIDIA/Megatron-LM#6697