[BUG] Delayed all_gather memory release using ZERO3

Describe the bug
I am using zero3 to train a model, the memory consumption is higher than expected, so I dumped a torch memory trace (below), during fwd and bwd, the all-gathered weights are not released immediately but delayed for several blocks, resulting in a higher memory peak.

I tried export TORCH_NCCL_AVOID_RECORD_STREAMS=1 but seems no difference.

To Reproduce
deepspeed: 0.15.0, 0.16.4 (both version tested, 0.16.4 is even worse)
torch: 2.6.0
cuda: 12.2

run the following code with 8gpus can reproduce:

import torch
import torch.distributed as dist
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration, AutoModelForCausalLM
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import StateDictType
from torch.distributed.fsdp.api import BackwardPrefetch, ShardingStrategy
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy, size_based_auto_wrap_policy


from torch.distributed.device_mesh import init_device_mesh
from transformers import AutoTokenizer, AutoModel
import os
import functools




os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"


def setup_distributed() -> None:
    """Initialize distributed training environment."""

    # local_rank = int(os.environ["LOCAL_RANK"])

    # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
    try:
        rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
    except KeyError as e:
        raise RuntimeError(f"Could not find {e} in the torch environment")

    local_rank = rank % 8
    torch.cuda.set_device(local_rank)

    # initialize the default process group
    dist.init_process_group(
        rank=rank,
        world_size=world_size,
        backend="nccl",
    )

def deepspeed_wrap(model):
    import deepspeed
    from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam

    optimizer = DeepSpeedCPUAdam(model.parameters(), lr=0.01)

    # DeepSpeed configuration (can be more complex)
    config_params = {
        "train_batch_size": 8,
        "gradient_accumulation_steps": 1,
        "zero_optimization":{
            "stage": 3,
            "offload_optimizer": {
                "device": "cpu",
                "pin_memory": True,
            },
        }
        ,
        "fp16": {
            "enabled": False
        }
    }

    # Initialize DeepSpeed engine
    model_engine, optimizer, _, _ = deepspeed.initialize(
        model=model,
        optimizer=optimizer,
        config_params=config_params
    )

    return model_engine

setup_distributed()
rank=dist.get_rank()
ws=dist.get_world_size()
print(f"{ws=}", flush=True)

pretrain_name_or_path="Qwen/Qwen2.5-7B"

model = AutoModelForCausalLM.from_pretrained(pretrain_name_or_path)
model.train()
model.gradient_checkpointing_enable()

model_engine = deepspeed_wrap(model)

input_ids = torch.randint(0, 10240, (1, 512), device='cuda')

for ind in range(5):
    if ind==1:
        torch.cuda.memory._record_memory_history()

    out = model_engine(input_ids=input_ids)["logits"]


    loss = out.mean()

    model_engine.backward(loss)

    model_engine.step()

    if ind==1:
        torch.cuda.memory._dump_snapshot(f"ds_train_7B_rank{rank}")

Start by running the provided Python reproduction with DeepSpeed ZeRO-3 on 8 GPUs and inspect the generated ds_train_7B_rank snapshot using the PyTorch memory trace tools. Compare when all-gathered weights are released during forward and backward passes. Done means identifying the cause of the delayed release and verifying reduced peak memory, but the issue names no project file or test to target.