No backward all-gather in deepseek simplefsdp+ep/tp

[ruisizhang123]

Bug description

After enabling SimpleFSDP on DeepSeek models, I found the backward all-gather used for reshard-after-forward behavior is missing. This happens to SimpleFSDP+ep/tp, but SimpleFSDP standalone has backward all-gathers. More details can be found in this pr description: #1529

In SimpleFSDP, this reshard-after-forward is achieved using checkpoint. probably sth is wrong with checkpoint API in SimpleFSDP+DeepSeek.

In eager-mode, there is no backward all-gather. In compile-mode, there will be all-gathers in tlparse, but I get the error as follows at the end of training.

[rank0]:[W808 15:00:07.207032687 ProcessGroup.cpp:364] Warning: At the time of process termination, there are still 20 unwaited collective calls. Please review your program to ensure that:
[rank0]:1. c10d_functional.wait_tensor() is invoked on all tensors returned from c10d_functional collective,
[rank0]:2. c10d_functional.wait_tensor() is invoked on all output tensors of async_op=True torch.distributed collective called under `with allow_inflight_collective_as_graph_input_ctx():`,
[rank0]:before the output tensors of the collective are used. (function ~WorkRegistry)

cc. @tianyu-l @anijain2305 @wconstab

Versions

See above

[tianyu-l]

cc @soulitzer seems one more issue around AC, lol

[ruisizhang123]

Some preliminary debugging progress:

In the compiled SimpleFSDP+EP trace, there are 18 all-to-all comm (but there should only be 8 all-to-all in debug model)...... The bwd FSDP in experts are first reduce-scattered and then all-gathered. These all look quite suspicious to me.

I also tried to print out all of the buffer in inductor, and found out all of the collective comm has its corresponding waits hmmmm.

[soulitzer]

I looked into the following warning a bit

[rank0]:[W808 15:00:07.207032687 ProcessGroup.cpp:364] Warning: At the time of process termination, there are still 20 unwaited collective calls. Please review your program to ensure that:
[rank0]:1. c10d_functional.wait_tensor() is invoked on all tensors returned from c10d_functional collective,
[rank0]:2. c10d_functional.wait_tensor() is invoked on all output tensors of async_op=True torch.distributed collective called under `with allow_inflight_collective_as_graph_input_ctx():`,
[rank0]:before the output tensors of the collective are used. (function ~WorkRegistry)

In my build, I was able to see 2 missing waits per iteration, but after getting rid of all graph breaks, there become only 1.

This single missing wait belongs to this all_reduce:

torch.ops._c10d_functional.all_reduce_.default(buf289, 'avg', '0')
buf292 = reinterpret_tensor(buf282, (1, 512, 256), (131072, 1, 512), 0); del buf282  # reuse
# Topologically Sorted Source Nodes: [], Original ATen: [_c10d_functional.wait_tensor]
stream0 = get_raw_stream(0)
triton_poi_fused_wait_tensor_9.run(buf289, buf292, 512, 256, stream=stream0)
# Topologically Sorted Source Nodes: [], Original ATen: [_c10d_functional.wait_tensor]
torch.ops._c10d_functional.wait_tensor.default(buf292)

(see _0_0_0/inductor_output_code_czrmcowzh2lp4sio64gimli4eu4aa552y5tswg6kih6bk5tuj5io_17.txt in tlparse)

In the above snippet, triton_poi_fused_wait_tensor_9 doesn't actually do any waiting! Instead it copies the contents of buf289 into a different buffer buf292

click to see triton_poi_fused_wait_tensor_9

# kernel path: /tmp/torchinductor_jw3468/jz/cjzslovjfzv6yw2vvc5qopgzng64jtw43bsztlv2me24axa57ye5.py
# Topologically Sorted Source Nodes: [], Original ATen: [_c10d_functional.wait_tensor]
# Source node to ATen node mapping:
# Graph fragment:
#   %buf32 : Tensor  = PlaceHolder[target=buf32]
#   %wait_tensor_37 : Tensor "f32[1, 512, 256][131072, 1, 512]cuda:0"[num_users=1] = call_function[target=torch.ops._c10d_functional.wait_tensor.default](args = (%all_reduce,), kwargs = {})
#   return %wait_tensor_37
triton_poi_fused_wait_tensor_9 = async_compile.triton('triton_poi_fused_wait_tensor_9', '''
import triton
import triton.language as tl

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
triton_helpers.set_driver_to_gpu()

@triton_heuristics.pointwise(
    size_hints={'y': 512, 'x': 256}, tile_hint=TileHint.SQUARE,
    filename=__file__,
    triton_meta={'signature': {'in_ptr0': '*fp32', 'out_ptr0': '*fp32', 'ynumel': 'i32', 'xnumel': 'i32', 'YBLOCK': 'constexpr', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=108, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
    inductor_meta={'grid_type': 'Grid2D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_wait_tensor_9', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 1, 'num_reduction': 0, 'backend_hash': 'F89C12E851B303B7D697BEBD762BD1121895275D7F92035C7DA8BE2DAB283256', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'tiling_scores': {'y': 1048576, 'x': 0}},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused_wait_tensor_9(in_ptr0, out_ptr0, ynumel, xnumel, YBLOCK : tl.constexpr, XBLOCK : tl.constexpr):
    ynumel = 512
    xnumel = 256
    yoffset = tl.program_id(1) * YBLOCK
    yindex = yoffset + tl.arange(0, YBLOCK)[:, None]
    ymask = yindex < ynumel
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[None, :]
    xmask = xindex < xnumel
    x1 = xindex
    y0 = yindex
    tmp0 = tl.load(in_ptr0 + (x1 + 256*y0), xmask & ymask, eviction_policy='evict_last')
    tl.store(out_ptr0 + (y0 + 512*x1), tmp0, xmask & ymask)
''', device_str='cuda')

Then, we wait on buf292 instead, so buf289 is never waited on, resulting in 512k being leaked per iteration.

Confusingly, although we only miss a single wait per iteartion, the same construction actually occurs 6 times total in codegened code. What is happening is for the first 5 times, the buf that doesn't get waited on is reinterpreted and reused later where it coincidentally gets waited on.

[xmfan]

cc @eellison who's looking at comms optimization passes. It seems strange for buf289 to be used before its wait_tensor op, Ruisi confirmed that there's no reordering is happening for this repro.

[ruisizhang123]

I looked into the following warning a bit

[rank0]:[W808 15:00:07.207032687 ProcessGroup.cpp:364] Warning: At the time of process termination, there are still 20 unwaited collective calls. Please review your program to ensure that:
[rank0]:1. c10d_functional.wait_tensor() is invoked on all tensors returned from c10d_functional collective,
[rank0]:2. c10d_functional.wait_tensor() is invoked on all output tensors of async_op=True torch.distributed collective called under `with allow_inflight_collective_as_graph_input_ctx():`,
[rank0]:before the output tensors of the collective are used. (function ~WorkRegistry)

In my build, I was able to see 2 missing waits per iteration, but after getting rid of all graph breaks, there become only 1.

This single missing wait belongs to this all_reduce:

torch.ops._c10d_functional.all_reduce_.default(buf289, 'avg', '0')
buf292 = reinterpret_tensor(buf282, (1, 512, 256), (131072, 1, 512), 0); del buf282  # reuse
# Topologically Sorted Source Nodes: [], Original ATen: [_c10d_functional.wait_tensor]
stream0 = get_raw_stream(0)
triton_poi_fused_wait_tensor_9.run(buf289, buf292, 512, 256, stream=stream0)
# Topologically Sorted Source Nodes: [], Original ATen: [_c10d_functional.wait_tensor]
torch.ops._c10d_functional.wait_tensor.default(buf292)

(see _0_0_0/inductor_output_code_czrmcowzh2lp4sio64gimli4eu4aa552y5tswg6kih6bk5tuj5io_17.txt in tlparse)

In the above snippet, triton_poi_fused_wait_tensor_9 doesn't actually do any waiting! Instead it copies the contents of buf289 into a different buffer buf292

click to see triton_poi_fused_wait_tensor_9
Then, we wait on buf292 instead, so buf289 is never waited on, resulting in 512k being leaked per iteration.

Confusingly, although we only miss a single wait per iteartion, the same construction actually occurs 6 times total in codegened code. What is happening is for the first 5 times, the buf that doesn't get waited on is reinterpreted and reused later where it coincidentally gets waited on.

Thank you for providing the workaround in moe.py to avoid graph graph. The "unwaited collective" problem is fixed in SimpleFSDP + EP (tlparse).

However, the communication pattern still looks a bit suspicious to me:

• In eager mode, (i) the backward all gather is still missing Profiler Trace. This is as expected because we only fixed all-to-all graph break in compile mode. (ii) the all-to-all disappeared 😭. I'm a bit uncertain if this is a SimpleFSDP problem or deepseek in torchtitan problem. Because with previous FSDP2+EP implementation, the all-to-all is also missing here. After pulling the latest main, FSDP2+EP has some unexpected bugs. (cc. @tianyu-l )

• In compile mode, (i) the backward all-to-all disappeared. It's a bit hard to narrow down if this is a compile problem or FSDP+EP implementation problem (because of eager-(ii)). (ii) in routed experts FSDP, the bwd all-gather & reduce scatter are missing. Profiler Trace

To Summarize:
(1) I also tried to narrow down the problem by changing checkpoint(self.replicate_compute, x, use_reentrant=False, context_fn=fsdp_policy) to self.replicate_compute(x), which basically remove reshard_after_forward in SimpleFSDP. In this profiler trace, compile mode's routed expert's bwd all-gather & reduce-scatter are still missing. Thus, I think it should be more of a SimpleFSDP + EP + torch.compile problem.

(2) The missing bwd all-gather in eager mode should still be a checkpoint + SimpleFSDP + EP problem lolll.

cc @anijain2305

[soulitzer]

(2) The missing bwd all-gather in eager mode should still be a checkpoint + SimpleFSDP + EP problem lolll.

It looks like SimpleFSDP's AC is applied to just the parameterization function (e.g. replicate_compute). This doesn't work because you want the all_gather to be triggered at the start of the backward of a FSDP region. To implement SimpleFSDP with AC, you need AC wrapping the entire FSDP region, e.g. FDSP(AC(...)).

This does, however, work in compile because AC in compile just works by tagging the corresponding nodes with "must recompute" for the partitioner to process.