PyTorch FSDP训练时显存不降反升，是配置问题还是预期行为？

这问题我踩过坑，单卡跑FSDP本来就会这样，分片省的是多卡间的显存，单卡上激活值和临时buffer反而可能更多。你试试把cpu_offload打开，或者调小bucket_cap_mb到25左右，显存能掉不少。另外LoRA的话，FSDP对冻结参数的处理有时会额外复制权重，检查下是不是把auto_wrap_policy设得太激进了。

FSDP的SHARD_GRAD_OP本身就只分片梯度，参数和优化器状态还是每卡一份的，7B模型光参数就要14G左右，加上LoRA的额外开销和激活值，70多G其实不算离谱。你对比DDP时可能没算上优化器状态，AdamW一开就是参数量的两倍显存，FSDP单卡反而要把整个模型的梯度都攒一轮再分片，峰值自然更高。建议你直接看下torch.profiler的内存快照，确认是参数、梯度还是激活占大头，另外试试把activation checkpointing打开，那个对显存影响比prefetch大得多。

这个现象我之前也踩过坑，SHARD_GRAD_OP其实只分片梯度，参数和优化器状态还是全量驻留的，7B模型光是参数+梯度+Adam状态就差不多要60G了，再加上LoRA的激活值，70多G不奇怪。你如果想让单卡显存真的降下来，得用FULL_SHARD，但代价是通信量翻倍，小batch下可能反而更慢。另外forward_prefetch对单卡场景基本没帮助，它主要是为了跨卡流水线并行设计的，你不如把cpu_offload打开试试，LoRA微调场景下能省不少。

看到你这个情况我第一反应是sharding_strategy设成SHARD_GRAD_OP其实只分片了梯度，参数和优化器状态还是全量复制，所以显存比DDP高完全正常，尤其是LoRA这种本身参数就少的情况，分片收益根本抵不过FSDP自身的通信缓冲和activation开销。你把策略改成FULL_SHARD试试，理论上参数和优化器状态都分片后应该能明显降下来，但代价是通信量会大不少，小batch下不一定划算。另外你用的是LoRA，其实可以考虑干脆不开FSDP，直接DDP配gradient checkpointing，7B在80G上完全能跑，还省去一堆调参的麻烦。

刚踩过类似的坑，试试关掉activation checkpointing或者调小batch size，这玩意分片后通信开销反而会放大峰值显存。

说实话你这情况我踩过一模一样的坑，最后发现大概率是预期行为。SHARD_GRAD_OP只分片了梯度和优化器状态，参数本身还是每卡都留一份完整副本，所以前向激活值加上参数副本，显存自然比DDP还高，尤其是7B这种规模，光参数就占14G多，LoRA虽然省了训练参数但基座权重没法省。另外你检查过activation_checkpointing开了没？FSDP的显存收益很大程度靠它把中间激活换到CPU，不然大batch下激活值才是吞显存的大头，比参数分片影响还猛。还有个小坑是forward_prefetch如果设置的时机不对，反而会提前把下一层的权重搬进来，导致峰值显存更难看，你可以试试把limit_all_gathers=True加上，强制延迟同步，我这边效果比调prefetch明显。最后建议你盯着torch.cuda.max_memory_allocated()看峰值，别只看nvidia-smi的实时占用，因为FSDP的all-gather和释放有滞后性，那个数字会误导人。如果还想压，就把sharding_strategy换成FULL_SHARD配合auto_wrap_policy，代价是通信开销大一点，但显存能掉到40G左右。

你这配置看起来没啥问题，但SHARD_GRAD_OP本来就是只分片梯度，参数和优化器状态还在每张卡上全量复制，7B模型光参数就占14G（半精度），加上LoRA的梯度和激活值，70多G其实不奇怪。想省显存得用FULL_SHARD，但代价是通信开销大，7B单卡跑本来就不太适合FSDP，分片收益要跨节点才明显。你DDP能跑的话，不如先查下是不是activation checkpointing没开，那个对显存影响比FSDP配置大得多。

大概率是FSDP把优化器状态也放在第一张卡上做初始化了，试试cpu_offload或者把batch size调小点看曲线。

说实话看到你这个现象我第一反应是检查一下你是不是把activation checkpointing和FSDP一起用了，因为光分片参数和梯度的话，7B模型在A100上激活值才是显存大头。SHARD_GRAD_OP这个策略下，前向传播时每层参数还是要全量物化到当前设备上的，如果你没有配合activate_checkpointing，激活值累积起来轻松超过DDP的占用。另外LoRA的话，base model的权重是冻结的，但FSDP的sharding单位是layer还是整个model会影响内存峰值，你可以试试把auto_wrap_policy设成按transformer block来分片，有时候能缓解瞬时峰值。还有个小坑是forward_prefetch在数据加载不均匀时反而会提前拉取下一层参数，导致当前层还没释放又占一份内存，你可以试着关掉它然后把limit_all_gathers=True打开，这个对峰值控制挺明显的。最后问一下，你确认过是纯参数+梯度+优化器状态的统计吗，因为PyTorch profiler里看起来的“显存”有时候会包含缓存分配器没释放的碎片，用torch.cuda.reset_peak_memory_stats()再测一次更准。我上次跑13B也遇到过类似情况，最后发现是mixed precision下master weight在每张卡上复制了一份，你要是开了fp16记得把cast_forward_inputs也检查下。

单卡跑FSDP本来就容易这样，分片没省多少反而多了all-gather开销，试试开CPU offload或者确认下是不是没启用LoRA的requires_grad。

单卡跑FSDP其实不太能体现分片优势，SHARD_GRAD_OP只切梯度，参数和优化器状态还是全量复制，显存自然下不来。7B模型光参数就14G，加上LoRA的激活值和临时buffer，70G挺正常的。想省显存可以试试FULL_SHARD配合CPU offload，或者干脆用QLoRA量化。另外prefetch那几个参数主要影响通信重叠，对峰值显存帮助有限，别指望靠它降。

显存涨这么多不太正常，检查下是不是没包住auto_wrap，或者LoRA层没被FSDP管到，参数其实还是全量复制了。

LoRA微调7B的话，SHARD_GRAD_OP其实没太大必要，因为LoRA的可训练参数本来就少，分片收益有限，反而FSDP的all-gather通信和临时buffer会额外吃显存。你可以试试FULL_SHARD配合把LoRA层单独排除在外，或者干脆用DDP加gradient checkpointing，显存可能更省。另外显存飙高也可能是optimizer state没被正确分片，检查下use_orig_params有没有开。

单卡A100跑FSDP确实有点尴尬，SHARD_GRAD_OP只分片梯度和优化器状态，参数还是全量复制的，7B模型光参数就占不少，加上LoRA的激活值和临时buffer，70G其实不算离谱。你试试FULL_SHARD，参数也切了显存会明显降，但通信开销更大。另外检查下auto_wrap_policy有没有正确套到Transformer层上，没wrap的话等于白配。

单卡跑FSDP确实容易出现这种情况，因为SHARD_GRAD_OP只分片梯度和优化器状态，参数本身还是全量复制的，加上FSDP自己的all-gather和reduce-scatter通信buffer，显存反而比DDP多。你要真想省显存得用FULL_SHARD，但那对单卡没啥意义，分片了也是自己扛全部。另外LoRA本身可训练参数很少，优化器状态占比不大，FSDP的收益就更不明显了，反而通信开销和临时buffer把显存吃掉了。建议单卡还是老老实实用DDP或者直接普通训练就行。

单卡跑FSDP确实容易踩这个坑，显存不降反升大概率不是配置写错了，而是FSDP本身的开销在单卡场景下被放大了。你用的SHARD_GRAD_OP只分片梯度和优化器状态，参数本身还是全量复制的，所以模型权重那份显存一点没省，反而多了all-gather通信时的临时缓冲区。7B的LoRA虽然可训练参数少，但基座模型的fp16权重就要占14G左右，加上all-gather的峰值和之前DDP的梯度桶，冲到70G不算离谱。forward_prefetch和backward_prefetch那几档在小规模或者单卡下基本看不出区别，它们的收益要等多卡、大batch、通信和计算能真正overlap的时候才体现出来。真要单卡省显存，不如直接上FULL_SHARD，把参数也切了，或者干脆用QLoRA那种4bit量化基座的方案，比调FSDP的prefetch实在得多。顺便问一句，你的A100是80G还是40G？如果是40G那这个数字就更值得警惕了。

单卡跑FSDP确实容易踩这个坑，你看到显存涨到70多G我一点都不意外。FSDP的SHARD_GRAD_OP模式只分片梯度和优化器状态，参数本身还是每张卡一份完整的，所以模型权重那部分显存根本没省下来。7B的LoRA虽然可训练参数少，但基座模型的fp32权重加上all-gather时的临时buffer，再加上FSDP为了通信效率会预取下一层参数，这些加起来很容易超过DDP。你要是单卡的话其实FSDP没啥收益，它本来就是为多卡设计的，单卡反而多了分片和all-gather的开销。forward_prefetch和backward_prefetch在单卡场景下确实没区别，因为压根没有跨卡通信，那几个选项就是给多卡调优用的。建议你确认下是不是真的只有一张卡在跑，如果是的话直接上DDP或者干脆不用分布式，LoRA微调7B单卡80G用普通训练循环完全够。如果确实是多卡环境，可以试试FULL_SHARD把参数也切了，但那样通信量会更大，得权衡一下。

FSDP显存比DDP高其实挺常见的，尤其在SHARD_GRAD_OP下，它只分片梯度和优化器状态，参数本身还是每个rank一份完整拷贝。7B模型fp16权重就14G左右，加上LoRA的额外参数、激活值、临时buffer，再算上all-gather时通信buffer的峰值，70多G不算离谱。你可能把FSDP想成了ZeRO-3那种全分片，但SHARD_GRAD_OP对应的是ZeRO-2，省的主要是优化器状态那块，对单卡权重占用帮助有限。想真正把参数也切走，得用FULL_SHARD，不过通信开销会明显上来。另外forward_prefetch和backward_prefetch在小规模或单节点内差异确实不大，它们主要影响的是通信和计算重叠的时机，显存峰值更多取决于all-gather的调度和临时buffer复用。你可以试试开activation_checkpointing，再把LoRA的target module限制一下，很多时候激活值才是隐形大户。还有个坑是如果用了auto_wrap_policy没包好，FSDP可能退化成一个巨大的flat parameter，分片效果会打折扣。

单卡跑FSDP其实挺容易踩坑的，SHARD_GRAD_OP只分片梯度，参数和优化器状态还是全量复制，显存自然比DDP高。7B模型光参数fp16就14G，加上LoRA的激活值和临时buffer，70G不算离谱。建议试试FULL_SHARD配合CPU offload，或者确认下auto_wrap_policy有没有正确切到Transformer层，不然分片粒度太粗等于没分。

LoRA微调场景下其实FSDP不一定划算，因为可训练参数本来就少，分片收益有限，反而FSDP自己的all-gather和reduce-scatter会额外吃显存，尤其是SHARD_GRAD_OP下前向还是要拉全参数。显存飙到70G我怀疑是模型加载时先完整实例化再加分片，可以试试用meta device加sync_module_states那套初始化流程，或者直接上QLoRA把基座压到4bit。另外7B单卡80G用DDP加LoRA其实是够的，FSDP在这时候更多是拖后腿。