# Not seeing memory benefit to accelerate/FSDP2

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fnot-seeing-memory-benefit-to-accelerate-fsdp2%2F158039%3E%5C
**Category:** 🤗Accelerate\
**Created:** [June 4, 2025, 9:34pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fnot-seeing-memory-benefit-to-accelerate-fsdp2%2F158039 "2025-06-04T21:34:41Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![hpcpony](https://avatars.discourse-cdn.com/v4/letter/h/779978/32.png) [@hpcpony](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fhpcpony)%5C
**Post date:** [June 4, 2025, 9:34pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fnot-seeing-memory-benefit-to-accelerate-fsdp2%2F158039%2F1 "2025-06-04T21:34:41Z")

</div>

TL;DR Why doesn’t Acclerate/FSDP seem to be doing much of anything to reduce memory in the following?

I’m trying to get some hands-on and learn how to run large models across multiple nodes and/or GPUs. I’m starting with Trainer/accelerate/FSDP2 and planning to work up from there but I think I’m missing something.

python 3.12.9  
torch 2.7.0  
transformers 4.52.4  
accelerate 1.7.0

My “toy” program to train an “empty” model:

```python
from datasets import Dataset, DatasetDict
from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM

from transformers import DefaultDataCollator, DataCollatorForLanguageModeling
from transformers import TrainingArguments, Trainer
import os

model_dir = 'NousResearch/Llama-3.2-1B'
TRACE = False
N = 2048
context_length = 64
batch_size = 64

def load_datasets() :
    train_data_list = [
        {"text" : "The quick brown fox jumped over the lazy dog's back t{:06d}".format(i)} for i in range(4*N)
        ]
    eval_data_list = [
        {"text" : "The quick brown fox jumped over the lazy dog's back e{:06d}".format(i)} for i in range(N)
        ]
    datasets = DatasetDict ( # create datasets dict train and eval
            { 'train': Dataset.from_list(train_data_list),
              'eval' : Dataset.from_list(eval_data_list)}
        )
    return datasets

def load_tokenizer(model_dir) :
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    return tokenizer

def load_model(model_dir) :
    # get just the config from the pretrained directory
    config = AutoConfig.from_pretrained(model_dir)
    model = AutoModelForCausalLM.from_config(config)
    return model

def mytrain(model_dir) :

    def tokenize(dataset) :
        return tokenizer(dataset['text'], padding='max_length', max_length=context_length, return_length=True)

    ##
    raw_datasets = load_datasets()
    if TRACE : print("dataset\n", raw_datasets)
    ##
    tokenizer = load_tokenizer(model_dir)
    if TRACE : print("tokenizer\n", tokenizer)
    ##
    tokenizer.pad_token = tokenizer.eos_token
    tokenized_datasets = raw_datasets.map(
        tokenize, batched=True, remove_columns=raw_datasets["train"].column_names)
    if TRACE : print("tokenized_datasets\n", tokenized_datasets)
    ##
    data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
    if TRACE :
        example_collated = data_collator([tokenized_datasets["train"][i] for i in range(3)])
        print("example_collated\n", example_collated)
    ##
    training_args = TrainingArguments( # do this before model load for FSDP?
        output_dir="outputs/",
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size,
        num_train_epochs=10,
        logging_strategy="epoch",
        eval_strategy="epoch",
        save_strategy="no",
        push_to_hub=False,
        disable_tqdm=True,
        deepspeed=None,
    )
    ##
    model = load_model(model_dir) # do the after TrainingArguments which sets up some stuff?
    if TRACE : print("model\n", model)
    ##
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_datasets["train"],
        eval_dataset=tokenized_datasets["eval"],
        processing_class=tokenizer,
        data_collator=data_collator,
    )
    trainer.train()

from datasets.utils.logging import disable_progress_bar
import torch
if __name__ == " __main__" :
  disable_progress_bar()
  mytrain(
     model_dir=model_dir
     )
  torch.distributed.destroy_process_group()

```

I first run my test progam as simple python/pytorch; single GPU without accelerate.

```shell
[gpu2:training] CUDA_VISIBLE_DEVICES=0 python 05_acctest.py 
{'loss': 0.8924, 'grad_norm': 0.8125, 'learning_rate': 4.50390625e-05, 'epoch': 1.0}
{'eval_loss': 2.5442957878112793, 'eval_runtime': 2.4496, 'eval_samples_per_second': 836.064, 'eval_steps_per_second': 13.063, 'epoch': 1.0}
{'loss': 0.6293, 'grad_norm': 0.65234375, 'learning_rate': 4.00390625e-05, 'epoch': 2.0}
{'eval_loss': 2.6600184440612793, 'eval_runtime': 2.4495, 'eval_samples_per_second': 836.094, 'eval_steps_per_second': 13.064, 'epoch': 2.0}
  .
  .
  .
{'loss': 0.6061, 'grad_norm': 0.4921875, 'learning_rate': 3.90625e-08, 'epoch': 10.0}
{'eval_loss': 2.8240463733673096, 'eval_runtime': 2.4496, 'eval_samples_per_second': 836.055, 'eval_steps_per_second': 13.063, 'epoch': 10.0}
{'train_runtime': 333.183, 'train_samples_per_second': 245.871, 'train_steps_per_second': 3.842, 'train_loss': 0.6405227959156037, 'epoch': 10.0}

```

While it’s running I use nvidia-smi to look at the memory used

```shell
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| 0 N/A N/A 21181 C python 21372MiB |
+-----------------------------------------------------------------------------------------+

```

That’s at least in the ball-park for what accelerate estimates:

```shell
[gpu2:training] accelerate estimate-memory NousResearch/Llama-3.2-1B
Loading pretrained config for `NousResearch/Llama-3.2-1B` from `transformers`...
┌────────────────────────────────────────────────────────┐
│ Memory Usage for loading `NousResearch/Llama-3.2-1B` │
├───────┬─────────────┬──────────┬───────────────────────┤
│ dtype │Largest Layer│Total Size│ Training using Adam │
├───────┼─────────────┼──────────┼───────────────────────┤
│float32│ 1002.0 MB │ 4.6 GB │ 18.42 GB │
│float16│ 501.0 MB │ 2.3 GB │ 9.21 GB │
│ int8 │ 250.5 MB │ 1.15 GB │ N/A │
│ int4 │ 125.25 MB │589.28 MB │ N/A │
└───────┴─────────────┴──────────┴───────────────────────┘

```

Next I use “accelerate config” to generate a config file for 2 GPUs using FSDP2. (mostly with default values)

```shell
[gpu2:training] cat 1n2gfsdp_defaults.yaml 
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
enable_cpu_affinity: false
fsdp_config:
  fsdp_activation_checkpointing: false
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false
  fsdp_reshard_after_forward: true
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
  fsdp_version: 2
machine_rank: 0
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

```

Using that file an running with accelerate…

```shell
[gpu2:training] CUDA_VISIBLE_DEVICES=0,1 accelerate launch --config_file 1n2gfsdp_defaults.yaml 05_acctest.py 
{'loss': 1.0797, 'grad_norm': 0.6328125, 'learning_rate': 4.5078125000000006e-05, 'epoch': 1.0}
{'eval_loss': 2.5193161964416504, 'eval_runtime': 1.376, 'eval_samples_per_second': 1488.383, 'eval_steps_per_second': 11.628, 'epoch': 1.0}
{'loss': 0.6584, 'grad_norm': 0.4609375, 'learning_rate': 4.0078125e-05, 'epoch': 2.0}
{'eval_loss': 2.5891079902648926, 'eval_runtime': 1.3771, 'eval_samples_per_second': 1487.218, 'eval_steps_per_second': 11.619, 'epoch': 2.0}
  .
  .
  .
{'loss': 0.6096, 'grad_norm': 0.462890625, 'learning_rate': 7.8125e-08, 'epoch': 10.0}
{'eval_loss': 2.754133462905884, 'eval_runtime': 1.3776, 'eval_samples_per_second': 1486.605, 'eval_steps_per_second': 11.614, 'epoch': 10.0}
{'train_runtime': 178.9799, 'train_samples_per_second': 457.705, 'train_steps_per_second': 3.576, 'train_loss': 0.6661747217178344, 'epoch': 10.0}

```

… nvidia-smi memory during the computation…

```shell
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| 0 N/A N/A 24421 C ...AI/training-4.52.4/bin/python 21384MiB |
| 1 N/A N/A 24422 C ...AI/training-4.52.4/bin/python 21388MiB |
+-----------------------------------------------------------------------------------------+

```

Next a config file with 4 GPUs…

```shell
[gpu2:training] cat 1n4gfsdp_defaults.yaml 
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
enable_cpu_affinity: false
fsdp_config:
  fsdp_activation_checkpointing: false
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false
  fsdp_reshard_after_forward: true
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
  fsdp_version: 2
machine_rank: 0
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

```

… execute using accelerate…

```shell
[gpu2:training] CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch --config_file 1n4gfsdp_defaults.yaml 05_acctest.py 
{'loss': 1.373, 'grad_norm': 0.458984375, 'learning_rate': 4.515625e-05, 'epoch': 1.0}
{'eval_loss': 2.402463912963867, 'eval_runtime': 0.6972, 'eval_samples_per_second': 2937.372, 'eval_steps_per_second': 11.474, 'epoch': 1.0}
{'loss': 0.7474, 'grad_norm': 0.435546875, 'learning_rate': 4.0156250000000004e-05, 'epoch': 2.0}
{'eval_loss': 2.3128156661987305, 'eval_runtime': 0.6946, 'eval_samples_per_second': 2948.607, 'eval_steps_per_second': 11.518, 'epoch': 2.0}
   .
   .
   .
{'loss': 0.6214, 'grad_norm': 0.30078125, 'learning_rate': 1.5625e-07, 'epoch': 10.0}
{'eval_loss': 2.432434320449829, 'eval_runtime': 0.694, 'eval_samples_per_second': 2950.801, 'eval_steps_per_second': 11.527, 'epoch': 10.0}
{'train_runtime': 89.6101, 'train_samples_per_second': 914.182, 'train_steps_per_second': 3.571, 'train_loss': 0.718875628709793, 'epoch': 10.0}

```

… nvidia-smi while executing…

```shell
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| 0 N/A N/A 25570 C ...AI/training-4.52.4/bin/python 20526MiB |
| 1 N/A N/A 25571 C ...AI/training-4.52.4/bin/python 20146MiB |
| 2 N/A N/A 25572 C ...AI/training-4.52.4/bin/python 20146MiB |
| 3 N/A N/A 25573 C ...AI/training-4.52.4/bin/python 20146MiB |
+-----------------------------------------------------------------------------------------+

```

Clearly something is happening; I’m getting a performance benefit from using more GPUs (almost linear!). But, I’m not seeing a substantial improvement in memory usage.

1. Is my config file missing something? Are there better parameters that facilitate memory savings?
2. Can I somehow get accelerate to dump what it thinks it’s doing (vs. what I specified in the config file)?
3. Can I somehow dump the wrapped model to see what FSDP has done?

===============================================================

I did a similar experiment with bloom-3b just to see if it made any difference, and things still seem strange.

```shell
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| 0 N/A N/A 37058 C python 74748MiB |
+-----------------------------------------------------------------------------------------+

┌────────────────────────────────────────────────────┐
│ Memory Usage for loading `bigscience/bloom-3b` │
├───────┬─────────────┬──────────┬───────────────────┤
│ dtype │Largest Layer│Total Size│Training using Adam│
├───────┼─────────────┼──────────┼───────────────────┤
│float32│ 2.39 GB │ 11.19 GB │ 44.74 GB │
│float16│ 1.2 GB │ 5.59 GB │ 22.37 GB │
│ int8 │ 612.5 MB │ 2.8 GB │ N/A │
│ int4 │ 306.25 MB │ 1.4 GB │ N/A │
└───────┴─────────────┴──────────┴───────────────────┘

+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| 0 N/A N/A 251138 C ...AI/training-4.52.4/bin/python 53922MiB |
| 1 N/A N/A 251139 C ...AI/training-4.52.4/bin/python 53538MiB |
| 2 N/A N/A 251140 C ...AI/training-4.52.4/bin/python 53538MiB |
| 3 N/A N/A 251141 C ...AI/training-4.52.4/bin/python 53538MiB |
+-----------------------------------------------------------------------------------------+

```

---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FJohn6666)%5C
**Post date:** [June 5, 2025, 6:24am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fnot-seeing-memory-benefit-to-accelerate-fsdp2%2F158039%2F2 "2025-06-05T06:24:05Z")

</div>

I don’t really understand how multi-GPU environments work…

> <https://github.com/pytorch/pytorch/issues/147168>
>
> 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)\`\](https://github.com/pytorch/pytorch/blob/20a369aa3abb6083600d5b22fcd8ba6e861c3959/torch/distributed/fsdp/\_fully\_shard/\_fully\_shard.py#L413) with explicit prefetch \`set\_modules\_to\_forward\_prefetch\` and \`set\_modules\_to\_backward\_prefetch\`.
> 
> \## Memory Over-Allocation
> 
> Then memory over-allocation happens like this:
> 
> !\[Image\](https://github.com/user-attachments/assets/058c2f0a-a15c-42f8-ac70-3bfa9f138008)
> 
> with memory traces:
> 
> !\[Image\](https://github.com/user-attachments/assets/34a43f32-7355-43ce-943b-59a074294de7)
> 
> !\[Image\](https://github.com/user-attachments/assets/c7122c74-a941-4e3a-b448-c15572b36a56)
> 
> 
> \## 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\](https://github.com/pytorch/pytorch/blob/0acbf8039abccfc17f9c8529d217209db5a7cc85/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L5373) (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\](https://pytorch.org/docs/stable/generated/torch.Tensor.record\_stream.html).)
> 
> 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\`\](https://github.com/pytorch/pytorch/blob/20a369aa3abb6083600d5b22fcd8ba6e861c3959/torch/distributed/fsdp/\_fully\_shard/\_fsdp\_param\_group.py#L92) 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\`.
> 
> \`\`\`python
> 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:
> \`\`\`python
> 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

> <https://github.com/pytorch/torchtune/issues/2402>
>
> Hi, when I set fsdp\_reshard\_after\_forward: False, the training speed increased b…y approximately 5-7%（tokens\_per\_second\_per\_gpu）. Are there any other configurations that affect performance? Or where do you recommend referring to for configurations?
> 
> In addition, the setting of gradient\_accumulation\_steps does not affect the speed. Generally speaking, setting a larger value will reduce the frequency of communication and speed up the training. The model used in the experiment is Qwen 2.5 3B.

> <https://github.com/pytorch/torchtitan/issues/735>
>
> \## Environment
> OS: Ubuntu
> GPU: 8x GPU
> torch: torch-2.6.0.dev20241212+cu124
> D…DP: 4-way Tensor Parallel \* 2-way FSDP
> 
> \## Problem
> I'm using FSDP+TP in my model and follow torchtitan code. when I switch fsdp1 to fsdp2, the memory usage showed by \`nvidia-smi\` increases by 10GB, also the peak active memory is greatly larger than fsdp1. is this expected? Which metric should be cared in \`memory\_summary\` to avoid OOM?
> 
> here is the result from \`torch.cuda.memory\_summary()\`. Following tables are generated when \*\*first step is end\*\*.
> 
> \* fsdp2
> \`\`\`
> |===========================================================================|
> | PyTorch CUDA memory summary, device ID 0 |
> |---------------------------------------------------------------------------|
> | CUDA OOMs: 0 | cudaMalloc retries: 0 |
> |===========================================================================|
> | Metric | Cur Usage | Peak Usage | Tot Alloc | Tot Freed |
> |---------------------------------------------------------------------------|
> | Allocated memory | 13975 MiB | 18803 MiB | 2142 GiB | 2128 GiB |
> | from large pool | 13959 MiB | 18790 MiB | 2140 GiB | 2127 GiB |
> | from small pool | 16 MiB | 17 MiB | 1 GiB | 1 GiB |
> |---------------------------------------------------------------------------|
> | Active memory | 13975 MiB | 39454 MiB | 2142 GiB | 2128 GiB |
> | from large pool | 13959 MiB | 39437 MiB | 2140 GiB | 2127 GiB |
> | from small pool | 16 MiB | 18 MiB | 1 GiB | 1 GiB |
> |---------------------------------------------------------------------------|
> | Requested memory | 13792 MiB | 39306 MiB | 2138 GiB | 2125 GiB |
> | from large pool | 13775 MiB | 39289 MiB | 2137 GiB | 2124 GiB |
> | from small pool | 16 MiB | 18 MiB | 1 GiB | 1 GiB |
> |---------------------------------------------------------------------------|
> | GPU reserved memory | 45590 MiB | 45590 MiB | 45590 MiB | 0 B |
> | from large pool | 45566 MiB | 45566 MiB | 45566 MiB | 0 B |
> | from small pool | 24 MiB | 24 MiB | 24 MiB | 0 B |
> |---------------------------------------------------------------------------|
> | Non-releasable memory | 377331 KiB | 7818 MiB | 1017 GiB | 1017 GiB |
> | from large pool | 375788 KiB | 7813 MiB | 1016 GiB | 1016 GiB |
> | from small pool | 1543 KiB | 10 MiB | 1 GiB | 1 GiB |
> |---------------------------------------------------------------------------|
> | Allocations | 4735 | 4738 | 34212 | 29477 |
> | from large pool | 1504 | 1507 | 15954 | 14450 |
> | from small pool | 3231 | 3348 | 18258 | 15027 |
> |---------------------------------------------------------------------------|
> | Active allocs | 4735 | 4738 | 34212 | 29477 |
> | from large pool | 1504 | 1507 | 15954 | 14450 |
> | from small pool | 3231 | 3348 | 18258 | 15027 |
> |---------------------------------------------------------------------------|
> | GPU reserved segments | 304 | 304 | 304 | 0 |
> | from large pool | 292 | 292 | 292 | 0 |
> | from small pool | 12 | 12 | 12 | 0 |
> |---------------------------------------------------------------------------|
> | Non-releasable allocs | 15 | 135 | 15054 | 15039 |
> | from large pool | 13 | 89 | 9160 | 9147 |
> | from small pool | 2 | 65 | 5894 | 5892 |
> |---------------------------------------------------------------------------|
> | Oversize allocations | 0 | 0 | 0 | 0 |
> |---------------------------------------------------------------------------|
> | Oversize GPU segments | 0 | 0 | 0 | 0 |
> |===========================================================================|
> \`\`\`
> 
> \* fsdp1
> \`\`\`
> |===========================================================================|
> | PyTorch CUDA memory summary, device ID 0 |
> |---------------------------------------------------------------------------|
> | CUDA OOMs: 0 | cudaMalloc retries: 0 |
> |===========================================================================|
> | Metric | Cur Usage | Peak Usage | Tot Alloc | Tot Freed |
> |---------------------------------------------------------------------------|
> | Allocated memory | 13947 MiB | 18561 MiB | 2156 GiB | 2142 GiB |
> | from large pool | 13937 MiB | 18556 MiB | 2155 GiB | 2141 GiB |
> | from small pool | 10 MiB | 11 MiB | 1 GiB | 1 GiB |
> |---------------------------------------------------------------------------|
> | Active memory | 13947 MiB | 25765 MiB | 2156 GiB | 2142 GiB |
> | from large pool | 13937 MiB | 25758 MiB | 2155 GiB | 2141 GiB |
> | from small pool | 10 MiB | 11 MiB | 1 GiB | 1 GiB |
> |---------------------------------------------------------------------------|
> | Requested memory | 13792 MiB | 25709 MiB | 2154 GiB | 2140 GiB |
> | from large pool | 13782 MiB | 25702 MiB | 2153 GiB | 2139 GiB |
> | from small pool | 9 MiB | 11 MiB | 1 GiB | 1 GiB |
> |---------------------------------------------------------------------------|
> | GPU reserved memory | 36458 MiB | 36458 MiB | 36458 MiB | 0 B |
> | from large pool | 36446 MiB | 36446 MiB | 36446 MiB | 0 B |
> | from small pool | 12 MiB | 12 MiB | 12 MiB | 0 B |
> |---------------------------------------------------------------------------|
> | Non-releasable memory | 402232 KiB | 6360 MiB | 1345 GiB | 1345 GiB |
> | from large pool | 400277 KiB | 6359 MiB | 1344 GiB | 1343 GiB |
> | from small pool | 1955 KiB | 6 MiB | 1 GiB | 1 GiB |
> |---------------------------------------------------------------------------|
> | Allocations | 2460 | 2463 | 26870 | 24410 |
> | from large pool | 832 | 835 | 14354 | 13522 |
> | from small pool | 1628 | 1631 | 12516 | 10888 |
> |---------------------------------------------------------------------------|
> | Active allocs | 2460 | 2463 | 26870 | 24410 |
> | from large pool | 832 | 835 | 14354 | 13522 |
> | from small pool | 1628 | 1631 | 12516 | 10888 |
> |---------------------------------------------------------------------------|
> | GPU reserved segments | 305 | 305 | 305 | 0 |
> | from large pool | 299 | 299 | 299 | 0 |
> | from small pool | 6 | 6 | 6 | 0 |
> |---------------------------------------------------------------------------|
> | Non-releasable allocs | 56 | 86 | 13297 | 13241 |
> | from large pool | 53 | 76 | 8544 | 8491 |
> | from small pool | 3 | 31 | 4753 | 4750 |
> |---------------------------------------------------------------------------|
> | Oversize allocations | 0 | 0 | 0 | 0 |
> |---------------------------------------------------------------------------|
> | Oversize GPU segments | 0 | 0 | 0 | 0 |
> |===========================================================================|
> \`\`\`
> 
> fsdp related code:
> \`\`\`Python
> compute\_dtype = torch.bfloat16 
> full\_shard: bool = True
> if use\_fsdp2:
> mixed\_fsdp2 = MixedPrecisionPolicy(reduce\_dtype=torch.float32, param\_dtype=compute\_dtype)
> for layer\_str, block in tp\_model.blocks.items():
> # fsdp2 currently don't change buffer dtype in mixed precision policy
> # so we have to set buffer dtype by hand
> block.t\_embed.to(torch.bfloat16)
> fully\_shard(block, mesh=ddp\_cp\_mesh, mp\_policy=mixed\_fsdp2)
> fully\_shard(tp\_model, mesh=ddp\_cp\_mesh, mp\_policy=mixed\_fsdp2)
> tp\_model\_ddp = tp\_model
> else:
> my\_auto\_wrap\_policy = functools.partial(
> transformer\_auto\_wrap\_policy,         
> transformer\_layer\_cls={
> type(tp\_model.blocks\["0"\]),
> },
> )
> st = ShardingStrategy.FULL\_SHARD if full\_shard else ShardingStrategy.SHARD\_GRAD\_OP
> mixed = MixedPrecision(param\_dtype=compute\_dtype, reduce\_dtype=torch.float32, buffer\_dtype=compute\_dtype)
> tp\_model\_ddp = FSDP(tp\_model, auto\_wrap\_policy=my\_auto\_wrap\_policy, device\_mesh=ddp\_cp\_mesh, mixed\_precision=mixed, 
> sharding\_strategy=st, device\_id=torch.cuda.current\_device(), use\_orig\_params=True)
> 
> \`\`\`
> \* fsdp2 memory timeline
> !\[image\](https://github.com/user-attachments/assets/729026bf-f630-47be-b9b6-31f1cdf2dd2f)
> 
> \* fsdp1 memory timeline
> !\[image\](https://github.com/user-attachments/assets/0fe45869-d28a-4a3f-821f-e03d836b7acf)

---

<div class="post-metadata">

**Author:** ![hpcpony](https://avatars.discourse-cdn.com/v4/letter/h/779978/32.png) [@hpcpony](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fhpcpony)%5C
**Post date:** [June 18, 2025, 3:49pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fnot-seeing-memory-benefit-to-accelerate-fsdp2%2F158039%2F3 "2025-06-18T15:49:22Z")

</div>

So after much futzing around and doing FSDP from pytorch I discovered that the answer to this question is that the memory usage reported by nvidia-smi is not an accurate reflection of memory required/used by pytorch. Apparently pytorch maintains a cache which is greater than that needed/used and that is primarily what the nvidia number reflects.

pytorch.cuda has a number of ways to get memory information that seems to be more relevant (though not always clear of the implications).

---

<div class="post-metadata">

**Author:** ![system](https://us1.discourse-cdn.com/hellohellohello/original/2X/d/de4155eb4aa4108ecb32a1389d7cc37ae69f88b7.png) [@system](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsystem)%5C
**Post date:** [June 19, 2025, 3:50am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fnot-seeing-memory-benefit-to-accelerate-fsdp2%2F158039%2F4 "2025-06-19T03:50:18Z")

</div>

Closed on June 19, 2025, 3:50am UTC

This topic was automatically closed 12 hours after the last reply. New replies are no longer allowed.