feat: Fix CPU offloading + add options for FSDP offload and expandable segments

[yfw]

What does this PR do ?

Addresses #33 (FSDP1) and #67 by:

1. Fixing the cpu offloading implementation so that HFPolicyWorker's memory is closer to 0 during VLLM generation.
2. Adding an option for using FSDP's built-in cpu offloading. With FSDP cpu offloading enabled, only the forward and backward passes are done on the GPU. Everything else, including the optimizer step update, is done on the CPU.
3. Adding the ability to set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True for the HFPolicy.

To test the impact of these settings, I did a sweep across different models, context lengths, and offload types. The results are here: https://docs.google.com/spreadsheets/d/1lWtw6-jbq4TAlM5Bu5Z3jNQjwAa_CEUX3bsKuMXHEkM/edit?usp=sharing

Offload Types:

• Main: Previous implementation on the main branch
• Manual: "Fixed" cpu offloading with the changes in this PR. This will be the default setting after this PR.
• FSDP: FSDP's built-in CPU offloading

Key Takeways:

• "Manual" offloading generally maintains or improves step time at lower allocated and reserved memory when compared with "Main"
• FSDP cpu offloading generally incurs some overhead in terms of Step 1 and Step 3 times when compared to "Manual" offloading.
  • An exception is for 8B, 7500 sequence length (row 23) which is the longest sequence length I tested without OOMing for "Manual" offloading
• Using expandable_segments keeps the reserved memory closer to allocated memory, but this is most evident with the "Main" offload type. Reserved memory is generally low for "Manual" and FSDP offload types.
  • One case where this does make a difference for the "Manual" offload type is the Llama3.1-8B model on 7500 context length, which OOMs without expandable_segments, but can run (slowly) without OOMing with expandable_segments (rows 22 and 23)
• For 8k context length and 8B model (FSDP1 memory usage still too high (8b on 8k seqlen not fitting) #67), FSDP cpu offloading is the only case that works without OOMing (Rows 30 and 31).
• Enabling expandable_segments also incurs some overhead in terms of step time, but this is most evident in the SFT case for 8k context, 8B model (rows 50-55), where expandable_segments runs ~3x slower than non-expandable_segments.

Issues

#33 (FSDP1 part)
#67

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[yfw]

Closing for #123