Add DTensor support for Ascend SwiGLU - #1448
Conversation
zheliuyu
left a comment
zheliuyu
left a comment
There was a problem hiding this comment.
Thanks for adding DTensor support on the Ascend SwiGLU path — the local-shard approach matches the CUDA SwiGLU pattern and looks right for an elementwise kernel. The comments below are for your reference.
| c_local = c_local * down_multiplier | ||
| ctx.save_for_backward(a_in, b_local) | ||
| ctx.dtensor_metadata = (device_mesh, placements) | ||
| return torch.distributed.tensor.DTensor.from_local(c_local, device_mesh, placements) |
There was a problem hiding this comment.
This DTensor path has no test delta. test_dtensor_liger_silumul is xfail when torch.cuda.device_count() < 8, so a 4×NPU host never actually runs it (world_size=8 cannot run there either). Please gate on infer_device() device count, keep importing LigerSiLUMulFunction from liger_kernel.ops so vendor replacement hits this class, and add at least one non-1.0 gate_multiplier / down_multiplier case — those scalars are applied in this wrapper, not in the Triton kernel.
There was a problem hiding this comment.
It looks like I marked the wrong line of code here, but it shouldn't affect the review comment. :>
There was a problem hiding this comment.
Thanks for catching this. I updated the DTensor test to gate each case using the device count from infer_device() and its world_size, while keeping LigerSiLUMulFunction imported from liger_kernel.ops. I also added coverage for non-default multipliers (0.7, 1.3). Verified on 4 Ascend NPUs: all four 4-rank cases passed, and the 8-rank cases were xfailed as expected.
| ctx.gate_multiplier = gate_multiplier | ||
| ctx.down_multiplier = down_multiplier | ||
|
|
||
| if isinstance(a, torch.distributed.tensor.DTensor) or isinstance(b, torch.distributed.tensor.DTensor): |
There was a problem hiding this comment.
Please bind torch.distributed.tensor at module import, the same way CUDA ops/swiglu.py and ops/rms_norm.py already do (try: import torch.distributed.tensor / from torch.distributed.tensor import DTensor as _DTensor, with a fallback). This isinstance(..., DTensor) runs on every Ascend SwiGLU forward; on torch 2.12+ the submodule is lazy, so this AttributeErrors even for plain tensors. torch_npu 2.9.0 does not catch that.
There was a problem hiding this comment.
I now bind DTensor and distribute_tensor at module import with a safe fallback, and use the bound symbols throughout the Ascend SwiGLU path.
Run elementwise SwiGLU on local DTensor shards and preserve their placements through forward and backward. Use HCCL as the Ascend distributed communication backend.
c1cf76e to
8550fa1
Compare
|
LGTM. cc @Tcc0403 Thanks for the revisions. |
Tcc0403
left a comment
Tcc0403
left a comment
There was a problem hiding this comment.
LGTM, thanks for the contribution
3 participants
Summary
Add DTensor support to the Ascend SwiGLU implementation.
SwiGLU now runs on local DTensor shards and preserves the original device
mesh and placements in forward and backward. Also fix
infer_comm_backend()to returnhcclon Ascend NPU.Testing Done
Environment: 4 x Ascend 910 NPU, CANN 9.1.0, PyTorch/torch-npu 2.9.0.
Non-DTensor SwiGLU tests:
67 passed, 18 skippedFour-device FP32/BF16 SwiGLU DTensor tests: passed
Four-device RMSNorm DTensor test for HCCL regression: passed
make checkstyle: passedHardware Type: 4 x Ascend 910 NPU
run
make testto ensure correctness (targeted tests run instead)run
make checkstyleto ensure code stylerun
make test-convergenceto ensure convergence (not applicable to this DTensor change)