Activation Offloading

[imh966]

Checklist Before Starting

• Search for similar PR(s).

What does this PR do?

This PR supports activation offloading, and currently it's only for FSDP backend.

High-Level Design

Our implementation is based on the one in TransformerEngine. For efficiency, it groups activations by TransformerLayer and offloads activation groups asynchronously. This means that the offloading of the i-th activation group and the computation of the i+1-th activation group happen at the same time, and there are at most two activation groups in GPU memory.

Specific Changes

1. Add activation offloading support.

API

Usage Example

export VLLM_ATTENTION_BACKEND=XFORMERS

python3 -m verl.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    data.train_files=./data/gsm8k/train.parquet \
    data.val_files=./data/gsm8k/test.parquet \
    data.train_batch_size=512 \
    data.max_prompt_length=512 \
    data.max_response_length=1024 \
    data.filter_overlong_prompts=True \
    data.truncation='error' \
    actor_rollout_ref.model.path=./huggingface.co/Qwen/Qwen2-7B-Instruct \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.model.use_remove_padding=True \
    actor_rollout_ref.actor.ppo_mini_batch_size=256 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=64 \
    actor_rollout_ref.actor.use_kl_loss=True \
    actor_rollout_ref.actor.kl_loss_coef=0.001 \
    actor_rollout_ref.actor.kl_loss_type=low_var_kl \
    actor_rollout_ref.actor.entropy_coeff=0 \
    actor_rollout_ref.model.enable_gradient_checkpointing=True \
    actor_rollout_ref.model.enable_activation_offload=True \
    actor_rollout_ref.actor.fsdp_config.param_offload=False \
    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=64 \
    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
    actor_rollout_ref.rollout.n=5 \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=64 \
    actor_rollout_ref.ref.fsdp_config.param_offload=True \
    algorithm.use_kl_in_reward=False \
    trainer.critic_warmup=0 \
    trainer.logger=['console','tensorboard'] \
    trainer.project_name='verl_grpo_example_gsm8k' \
    trainer.experiment_name='qwen2_7b_function_rm' \
    trainer.n_gpus_per_node=8 \
    trainer.val_before_train=False \
    trainer.nnodes=1 \
    trainer.save_freq=-1 \
    trainer.test_freq=5 \
    trainer.total_epochs=15

Test

We conducted experiments on the Qwen2 7B model based on the above script. The memory and throughput data are shown in the figures below, where the blue line represents activation offloading.

Additional Info.

• Issue Number: none
• Training: This PR will affect FSDP backend
• Inference: none

Checklist Before Submitting

• Read the Contribute Guide.
• Apply pre-commit checks.
• Add [BREAKING] to the PR title if it breaks any API.
• Update the documentation about your changes in the docs.
• Add CI test(s) if neccessary.

[CLAassistant]

All committers have signed the CLA.

[casper-hansen]

What makes the activation checkpointing incompatible with activation offloading? In most cases, you would like to combine both for maximum memory saved.

[imh966]

It's just that the activation checkpointing implementation in the transforms library does not work well with activation offloading directly. On the one hand, when activation checkpointing is enabled, Verl assigns "use_reentrant" to false, which leads to an implementation based on saved_tensor_hook. My implementation is also based on saved_tensor_hook, while PyTorch can only apply one saved_tensor_hook, unlike forward_hook, which allows multiple hooks. On the other hand, if try to assign "use_reentrant" to true, the forward function of the FSDP module will be executed again in backward, resulting in an additional allgather for the parameters of the FSDP module.

In my implementation, you can still combine them to save more memory. When activation offloading is enabled, I disable activation checkpointing in transformers library and enable it again in another way.

[ZihengJiang]

Remove the debug argument in the signature

[ZihengJiang]

not used?

[ZihengJiang]

Thanks for the PR. Took a glance and left some comments. Please add some tests to make sure that the feature works properly.

Also, it seems fsdp has some support for activation offloading: https://github.com/pytorch/pytorch/blob/main/torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py#L275

Have you checked this?

[BearBiscuit05]

#1118, This PR can fix the issue where vLLM doesn't detect memory release.

[eric-haibin-lin]

please add a unit test to CI

[eric-haibin-lin]

pls update PR description according to https://github.com/volcengine/verl/blob/main/.github/PULL_REQUEST_TEMPLATE.md and complete the checklist

[imh966]

Thanks for the PR. Took a glance and left some comments. Please add some tests to make sure that the feature works properly.

Also, it seems fsdp has some support for activation offloading: https://github.com/pytorch/pytorch/blob/main/torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py#L275

Have you checked this?

Thanks for the PR. Took a glance and left some comments. Please add some tests to make sure that the feature works properly.

Also, it seems fsdp has some support for activation offloading: https://github.com/pytorch/pytorch/blob/main/torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py#L275

Have you checked this?

Yes, I checked this support before the PR. It does not support to overlap offloading and computation, and it offloads some parameters of module.

[ZihengJiang]

@imh966 thanks for confirming. Please add a test in the CI

[ZihengJiang]

It will disable gradient checkpointing silently while activation offload is enabled in the same time, right? This sounds strange for user.

Let's just verify the configuration and raise an explicit error if enable_activation_offload and enable_gradient_checkpointing is enabled in the same time.

[ZihengJiang]

Also, given this feature only support FSDP for now, adding assertion in configuration verification to make sure `strategy == "fsdp".

[imh966]

Thank you for your suggestions. Actually, it just disables gradient checkpointing in transformers library and enables it in another way because gradient checkpointing in transformers library is incompatible with my implementation of activation offloading. I will try to add some comments to make it not so strange for user.

[ZihengJiang]

I suggest short name for this utilities function, just offload_activation(model) and add assert for FSDP only.

[imh966]

I've changed the name of this function to enable_activation_offloading. Do you think it's fine?

[ZihengJiang]

sounds good to me

[ZihengJiang]

@eric-haibin-lin @vermouth1992 shall we merge this?

[wplf]

@imh966 I've seen the offload implementation between yours and TEs.
Do you modify TE's offload to support FSDP?
Is there other features besides it?

[imh966]

@imh966 I've seen the offload implementation between yours and TEs. Do you modify TE's offload to support FSDP? Is there other features besides it?

@wplf Actually, it's almost a copy from TE. I just deleted some unused code to eliminate dependence on TE and added a few modifications to make it work well with FSDP.

[wplf]

@imh966 I've seen the offload implementation between yours and TEs. Do you modify TE's offload to support FSDP? Is there other features besides it?

@wplf Actually, it's almost a copy from TE. I just deleted some unused code to eliminate dependence on TE and added a few modifications to make it work well with FSDP.

Thank you for your reply. I am impressed very much!

[wplf]

Thank you for great work. I've test this pull request on my own FSDP repo. And it works very well.