[feat] support activation cpu offload in fsdp and fsdp2#7201
[feat] support activation cpu offload in fsdp and fsdp2#7201Jintao-Huang merged 7 commits intomodelscope:mainfrom
Conversation
Summary of ChangesHello @meichangsu1, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a critical memory optimization feature: activation CPU offloading for FSDP and FSDP2 training strategies. By intelligently moving activation tensors to CPU memory, the system can train larger models or use more extensive batch sizes on GPUs with limited memory. The implementation includes a sophisticated asynchronous offloading mechanism and is seamlessly integrated into the existing training framework via a new callback, making it easy to enable and configure. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for activation CPU offload in FSDP and FSDP2, which is a great feature for memory optimization. The implementation adds a new ActivationCpuOffloadCallBack and the associated logic to hook into the model's forward pass.
My review has identified a few issues:
- There is a critical inconsistency in the FSDP configuration key used for the FSDP version (
versionvs.fsdp_version), which will prevent the new feature from working. - The example configuration file has a confusing note.
- The example training script has a typo.
- The new
activation_cpu_offload.pyfile has some minor issues with logging configuration and type hinting.
I have provided detailed comments and suggestions to address these points. Once these are resolved, the PR should be in good shape.
| CUDA_VISIBLE_DEVICES=0,1 \ | ||
| swift sft \ | ||
| --model 'Qwen/Qwen3-0.6B' \ | ||
| --dataset 'swift/self-cognition#1000' \ \ |
There was a problem hiding this comment.
There is an extra backslash here (\\). A single backslash is sufficient for line continuation in shell scripts. This could cause unexpected behavior in some shells.
| --dataset 'swift/self-cognition#1000' \ \ | |
| --dataset 'swift/self-cognition#1000' \ |
| self.model_parameters_storage = new_storage | ||
|
|
||
|
|
||
| def get_torch_device() -> any: |
There was a problem hiding this comment.
The type hint any is used for the return value. any is a built-in function in Python. The correct type hint for an arbitrary type is Any from the typing module, which is already imported in this file.
| def get_torch_device() -> any: | |
| def get_torch_device() -> Any: |
|
方便问一下，这个RP解决的是什么问题啊 |
| @@ -0,0 +1,27 @@ | |||
| #!/bin/bash | |||
There was a problem hiding this comment.
加一个 前后的显存占用对比吧
然后用8B的模型
| from swift.utils import get_logger | ||
|
|
||
| logger = get_logger() | ||
| logger.setLevel(logging.WARNING) |
|
有参考代码的链接嘛 |
|
|
|
||
| # activation_cpu_offload=false | ||
| # OOM | ||
| # {'loss': 1.13790035, 'grad_norm': 1.41472316, 'learning_rate': 5e-05, 'token_acc': 0.83174487, 'epoch': 0.04, 'global_step/max_steps': '1/27', 'percentage': '3.70%', 'elapsed_time': '46s', 'remaining_time': '20m 1s', 'memory(GiB)': 61.79, 'train_speed(iter/s)': 0.021641} |
There was a problem hiding this comment.
不是这个例子跑的，用的内部一个数据集跑的，那个数据机的token 长度比较大
| --lora_rank 8 \ | ||
| --lora_alpha 32 \ | ||
| --target_modules all-linear \ | ||
| --freeze_vit true \ |
There was a problem hiding this comment.
sorry for that,the script displayed here is copied from the other demo;while the log is not produced from the train script here ，i will update it later
- Add ActivationCpuOffloadCallBack import and registration in callbacks mapping - Automatically append activation_cpu_offload callback when FSDP config has activation_cpu_offload enabled - Enables memory-efficient training by offloading activations to CPU during FSDP forward pass
…offload - Add `fsdp2.json` configuration file for PyTorch native FSDP v2 with activation CPU offloading - Include detailed parameter documentation and usage notes for FSDP2 - Provide example training script (`train.sh`) demonstrating multi-GPU training with LoRA - Disable gradient checkpointing in favor of FSDP's native activation checkpointing - Enable CPU RAM efficient loading and sharded state dicts for memory optimization
meichangsu1
force-pushed
the
activation_cpu_offlod
branch
from
86979cb to
20680af
Compare
- Add __init__ method to ActivationCpuOffloadCallBack to properly initialize parent class - Update import to use local base TrainerCallback instead of transformers version - Ensure callback follows consistent initialization pattern with other callbacks
- Remove activation_cpu_offload parameter from fsdp2.json - Set activation_checkpointing to true for improved memory efficiency - Maintain existing auto_wrap_policy and state_dict_type settings
| @@ -0,0 +1,26 @@ | |||
| { | |||
There was a problem hiding this comment.
放在examples/ascend文件夹下吧
| @@ -0,0 +1,54 @@ | |||
| #!/bin/bash | |||
| CUDA_VISIBLE_DEVICES=0,1 \ | |||
| NPROC_PER_NODE=2 \ | |||
There was a problem hiding this comment.
这能给个体现这个技术很好的例子嘛
因为看这里例子，显存降低的不多，但是速度降了很多
| self.model_parameters_storage = new_storage | ||
|
|
||
|
|
||
| def get_torch_device() -> any: |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a valuable feature for activation CPU offloading with FSDP, which can significantly reduce GPU memory usage. The implementation is thorough, including asynchronous offloading with double buffering and compatibility with activation checkpointing. I've identified a critical bug that prevents the feature from being enabled, along with a couple of other suggestions for improvement to enhance robustness and clarity. Overall, this is a great addition.
| def get_torch_device() -> any: | ||
| """Return the corresponding torch attribute based on the device type string. | ||
| Returns: | ||
| module: The corresponding torch device namespace, or torch.cuda if not found. | ||
| """ | ||
| device_name = get_device_name() | ||
| try: | ||
| return getattr(torch, device_name) | ||
| except AttributeError: | ||
| logger.warning(f"Device namespace '{device_name}' not found in torch, try to load torch.cuda.") | ||
| return torch.cuda |
There was a problem hiding this comment.
The get_torch_device function is not robust for CPU-only environments. If get_device_name() returns 'cpu', getattr(torch, 'cpu') raises an AttributeError, and the code falls back to torch.cuda. If CUDA is not available, the subsequent call to torch.cuda.Stream() will crash. Since this feature is designed for offloading from a GPU to save memory, it's better to explicitly raise an error in CPU-only mode to avoid unexpected crashes and provide a clear message to the user.
| def get_torch_device() -> any: | |
| """Return the corresponding torch attribute based on the device type string. | |
| Returns: | |
| module: The corresponding torch device namespace, or torch.cuda if not found. | |
| """ | |
| device_name = get_device_name() | |
| try: | |
| return getattr(torch, device_name) | |
| except AttributeError: | |
| logger.warning(f"Device namespace '{device_name}' not found in torch, try to load torch.cuda.") | |
| return torch.cuda | |
| def get_torch_device() -> any: | |
| """Return the corresponding torch attribute based on the device type string. | |
| Returns: | |
| module: The corresponding torch device namespace, or torch.cuda if not found. | |
| """ | |
| device_name = get_device_name() | |
| if device_name == 'cpu': | |
| raise RuntimeError('Activation CPU offload requires a device with streams (e.g., CUDA, NPU) and is not supported in CPU-only mode.') | |
| try: | |
| return getattr(torch, device_name) | |
| except AttributeError: | |
| logger.warning(f"Device namespace '{device_name}' not found in torch, try to load torch.cuda.") | |
| return torch.cuda |
- Add new training script `train.sh` for Ascend platform with activation CPU offload configuration - Include comprehensive training parameters for Qwen3-8B model with LoRA fine-tuning - Provide example training outputs for both activation_cpu_offload=true and false scenarios - Move existing fsdp2.json configuration file to Ascend examples directory
|
哈咯，lint 过一下 |
…c copy - Change training dataset from `AI-ModelScope/LongAlpaca-12k` to `AI-ModelScope/alpaca-gpt4-data-zh` in the example script - Modify `SynchronizedGroupOffloadHandler.offload` to use synchronous, non-pinned memory copy when NPU is available, as NPU does not fully support async H2D/D2H with pinned memory
|
done |
2 participants
[feat] support activation cpu offload in fsdp and fsdp2
PR type
PR information
Introduce activation CPU offloading built on autograd saved_tensors hooks and grouped async offload/reload to reduce GPU activation memory.
Add synchronous and async double-buffer offload handlers with stream-based D2H/H2D overlap and group window scheduling.
Wrap FSDP/FSDP2 layer forward to insert group commit boundaries and manage offload context; skip Embedding layers.
Provide activation checkpointing compatibility by replacing transformers’ checkpointing with internal checkpoint wrapper when enabled.
Add training callback to read fsdp_config flags and enable activation offload for FSDP v1/v2 at train start.
Experiment results
a simple test
for model-specific results,please look the examples in the code