[FSDP2] The evil `record_stream` in c10d causes FSDP2 to over-allocate GPU memory

[leonardo0lyj]

Hey Andrew @awgu,

As a big fan of FSDP2, I find an potential bug 😄

Demand:

• No inter-stream memory fragmentation (incurred by copy in streams)
• Explicit Prefetch
• CPU runs a head of GPU by a lot

_set_unshard_async_op(True)

To satisfy these demands, FSDP2 has to turn on _set_unshard_async_op(True) with explicit prefetch set_modules_to_forward_prefetch and set_modules_to_backward_prefetch.

Memory Over-Allocation

Then memory over-allocation happens like this:

with memory traces:

Root Cause

As known to all, these memory over-allocations are caused by the evil tensor.record_stream(ncclStream). Although FSDP2 tried to avoid this evil originated from FSDP1, such record_stream still is embedded in all c10d collectives (when async_op=True). Therefore, FSDP2 still suffers over-allocation from this evil in c10d.

Candidate Solution

Not sure how can we avoid the record_stream even when async_op=True?

IMO, candidate solutions are below:

1. Make TORCH_NCCL_AVOID_RECORD_STREAMS=True as an default value, getting rid of the record_stream in c10d. (Safety should be good without record_stream, as collective with async_op=True usually starts from allocation stream and ends at allocation stream, or users indeed know how to manually sync streams.)

2. Make TORCH_NCCL_AVOID_RECORD_STREAMS=True an advanced option to each collective, such as dist.all_gather(..., _avoid_record_stream=True). This limits the scope of environmental TORCH_NCCL_AVOID_RECORD_STREAMS to each specific collective.

3. Use only dist.all_gather(async_op=False) in FSDP2, but changes the current_stream to the all_gather_stream such that all gather still allocates/frees in current_stream while runs in all_gather_stream and overlaps with current_stream, just like async_op=True.

def get_all_gather_streams(
        self, async_op: bool, training_state: TrainingState
    ) -> tuple[torch.Stream, torch.Stream]:
        if not async_op and training_state in (
            TrainingState.FORWARD,
            TrainingState.PRE_BACKWARD,
        ):
            # Use separate streams for implicit prefetching
            return self.all_gather_copy_in_stream, self.all_gather_stream
        
        # Use separate streams for explicit prefetching!
        current_stream = self.device_handle.current_stream()
        return current_stream, self.all_gather_stream # Change this!

How do you prefer?

(Let us make FSDP great again 😄)

Code

P.S. the code to reproduce over-allocation:

class MLP(nn.Module):
    def __init__(self, hidden_dim: int, bias: bool = False):
        super().__init__()
        self.fc1 = nn.Linear(hidden_dim, hidden_dim, bias=bias)
        self.gelu = nn.GELU()
        self.fc2 = nn.Linear(hidden_dim, hidden_dim, bias=bias)

    def forward(self, x):
        x = self.fc1(x)
        x = self.gelu(x)
        x = self.fc2(x)
        return x


class MultiMLP(nn.Module):
    def __init__(self, hidden_dim: int, bias: bool = False, layers: int = 4):
        super().__init__()
        self.pre_norm = nn.LayerNorm(hidden_dim, bias=bias)
        self.mlps = nn.ModuleList([MLP(hidden_dim, bias) for _ in range(layers)])
        self.post_norm = nn.LayerNorm(hidden_dim, bias=bias)

    def forward(self, x):
        x = self.pre_norm(x)
        for mlp in self.mlps:
            x = x + mlp(x)
        x = self.post_norm(x)
        return x

class TestMemory(DTensorTestBase):
    @with_comms
    def test_over_allocation(self):
        mesh = init_device_mesh("cuda", (self.world_size,))
        device = torch.device("cuda")
        hidden_dim = 10240
        total_bsz = 16

        # ----- init model --------
        torch.manual_seed(0)
        model = MultiMLP(hidden_dim=hidden_dim).to(device).to(torch.float32)

        # --------  fsdp2 wrap --------
        fully_shard_fn = functools.partial(
            fully_shard,
            mesh=mesh,
            reshard_after_forward=True,
        )

        last_fsdp_module = None
        for module in model.modules():
            if isinstance(module, MLP):
                fully_shard_fn(module)
                if last_fsdp_module is not None:
                    last_fsdp_module.set_modules_to_forward_prefetch([module])
                    module.set_modules_to_backward_prefetch([last_fsdp_module])
                last_fsdp_module = module
        fsdp_model = fully_shard_fn(model)
        fsdp_model._set_unshard_async_op(True)

        optim = torch.optim.Adam(fsdp_model.parameters())

        # ----- init data -----
        torch.manual_seed(self.rank)
        bsz = total_bsz // self.world_size

        # --------  training loop --------
        torch.distributed.barrier()
        torch.cuda.synchronize(self.rank)
        
        train_iter = 4
        for iter in range(train_iter):
            # torch.distributed.barrier()
            # torch.cuda.synchronize(self.rank)

            if self.rank == 0 and iter == train_iter - 1:
                torch.cuda.memory._record_memory_history(max_entries=int(1E6))

            with record_function("## zero grad ##"):
                optim.zero_grad()

            input = torch.randn((bsz, hidden_dim), device="cuda")

            with record_function(f"## forward ##"):
                output = fsdp_model(input)
                loss = output.mean()

            with record_function(f"## backward ##"):
                loss.backward()

            with record_function("## optimizer step ##"):
                optim.step()

            if self.rank == 0 and iter == train_iter - 1:
                timestamp = datetime.now().strftime("%b_%d_%H_%M_%S")
                file_name = f"mem_{timestamp}"
                torch.cuda.memory._dump_snapshot(f"{file_name}.pickle")
                torch.cuda.memory._record_memory_history(enabled=None)

        torch.distributed.barrier()
        torch.cuda.synchronize(self.rank)

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @zhaojuanmao @mrshenli @rohan-varma @chauhang

[awgu]

Thanks for the detailed issue!

I think we should do TORCH_NCCL_AVOID_RECORD_STREAMS=True by default as the preferred solution.

[awgu]

cc: @wconstab do you think we can try to roll out TORCH_NCCL_AVOID_RECORD_STREAMS=True as default? The only cases where this could be considered a regression are when the user does not call .wait() on a collective on the CPU side before the GPU usage has finished. However, this is not the expected usage, and there are likely far more people being burned by memory spikes due to recordStream than those not calling .wait() when they expect their collective tensors to be free-able.

[wconstab]

Yea, we've been talking about doing this for a while, can folks help remind me why we didn't do it? @kwen2501 @shuqiangzhang

One note is that it is quite common to write p2p send ops without calling wait. The work obj keeps the tensor alive. Will we break these cases by changing the default?

[awgu]

TORCH_NCCL_AVOID_RECORD_STREAMS=True has no effect on P2P ops currently, so switching the default should have no impact on them.

[d4l3k]

Just to make sure I understand, the issue here is that the allocations persist too long? How does CUDACachingAllocator know when to free the memory? I'm not seeing much documentation explaining the behavior

[kwen2501]

cc: @wconstab do you think we can try to roll out TORCH_NCCL_AVOID_RECORD_STREAMS=True as default? The only cases where this could be considered a regression are when the user does not call .wait() on a collective on the CPU side before the GPU usage has finished. However, this is not the expected usage, and there are likely far more people being burned by memory spikes due to recordStream than those not calling .wait() when they expect their collective tensors to be free-able.

Yeah, this is the reason we did not turn it on by default when async_ops=True.

Given that FSDP2 added async=True and the pain described, can we consider a context manager that has same effect as setting TORCH_NCCL_AVOID_RECORD_STREAMS=True? This is more flexible -- giving control to FSDP.

Otherwise, if c10d sets it to True in general, I think we would need to educate our users that: if you don't call .wait(), c10d will hold your tensor as hostage. Not entirely impossible, but there is some BC smell in it.

[fegin]

@d4l3k @awgu wrote a nice explanation: https://github.com/pytorch/torchtitan/pull/737/files

[kwen2501]

if you don't call .wait(), c10d will hold your tensor as hostage

Giving this a second thought, maybe we can rely on watchdog to do the unstash if the user did not call wait()? As watchdog will check whether a collective has finished. In this case, the effect would not be much different from CudaCachingAllocator doing the release, as in the record_stream case. But at least it guarantees that the tensor will be released.

[awgu]

Otherwise, if c10d sets it to True in general, I think we would need to educate our users that: if you don't call .wait(), c10d will hold your tensor as hostage. Not entirely impossible, but there is some BC smell in it.

I agree that there is a BC issue, and this is what has been stopping people from flipping the default. However, I am currently proposing that breaking this BC is preferred.

Every training codebase I have seen that uses async_op=True collectives has to set TORCH_NCCL_AVOID_RECORD_STREAMS=1. If you can find some real examples where people prefer the =0 behavior, then please do share. I have not found any.

If there is a strong case to preserve the existing behavior, that is fine -- we can leave an option for it, but let us not make it the default. IMHO, the default behavior should not be a rough edge that many users are repeatedly being burned by. cc: @kwen2501

---

How does CUDACachingAllocator know when to free the memory? I'm not seeing much documentation explaining the behavior

When the user calls tensor.record_stream(), the allocator will record a CUDA event and query this event upon every subsequent tensor allocation until the event is completed, at which point the allocation can be freed. record_stream is the only case in PyTorch where the tensor's lifetime depends on GPU execution timing, not CPU issue order. cc: @d4l3k

[leonardo0lyj]

Support @awgu 👍

Every training codebase I have seen that uses async_op=True collectives

So true. To my personal experience (studying at training team of five institutions: ByteDance, Meta, Microsoft, IBM, academia) over the past 10 years, I recall most c10d usages are async_op=True, and every async_op=True is properly followed with wait(). As the official pytorch distributed doc says async_op=True + wait(), which is the entry-level note for users of c10d:

--

By contrast, the issue of memory hostage and record_stream is far too advanced for common users, and takes great effort to debug, especially it is non-deterministic.

As a lazy user, I just want to call async_op=True + wait(), without worrying if there is memory hostage in c10d, or did I turned on TORCH_NCCL_AVOID_RECORD_STREAMS, or did I read doc of TorchTitan.

[kwen2501]

Chatted with @awgu offline, here is a brief summary:

The watchdog-as-a-safety-net solution is purposed to remove the BC smell.
It is more of a "self-discipline" detail of c10d -- so we won't document this to the users.
We would instead document the expectation of wait() call after an async collective call. In which case, the safety net is a no-op.

With that, we will flip on TORCH_NCCL_AVOID_RECORD_STREAMS=True as default.

[leonardo0lyj]

Great news! Looking forward~ Thanks @kwen2501 and @awgu 😊