Enable PPO on Intel XPU using a tiny model

[songhappy]

Context

What is the purpose of this PR? Is it to

• add a new feature

Please link to any issues this PR addresses.
https://jira.devtools.intel.com/browse/IPB-2914

Changelog

Added a configuration file of running PPO on Intel PVC 48G.

Test plan

• manually run any new or modified recipes with sufficient proof of correctness
• I did not change any public API

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2446

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 09dd2a6 with merge base 0e8f840 ():

NEW FAILURE - The following job has failed:

• Lint / lint (3.10) (gh)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 23.13%. Comparing base (cf0142b) to head (09dd2a6).
Report is 13 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #2446       +/-   ##
===========================================
- Coverage   65.38%   23.13%   -42.26%     
===========================================
  Files         374      379        +5     
  Lines       22172    22793      +621     
===========================================
- Hits        14498     5273     -9225     
- Misses       7674    17520     +9846     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[SalmanMohammadi]

Hi @songhappy. Thanks for opening this. It's great to see support in torchtune for different hardware and smaller models, but I'm hesitant to land a config for an xpu device as we won't be able to test and maintain it going forward. We generally test and ship our configs on cuda and allow users to configure it on their own.

[ebsmothers]

Hi @songhappy thanks for the PR! For this case I think it would make sense to host this config outside of torchtune core. If I understand correctly the main difference is that this will run on XPU + the usage of a tiny Llama model, right? We don't really have any configs with tiny llama (as far as I know) so I think it would be a bit strange to add for just this one case. Let me know if this makes sense to you, thanks!

[songhappy]

Thanks for reviewing it. Actually, TinyLlama is not really tiny, it's 1B, and we have a couple of 1B or 0.5B model configurations in here. For example, https://github.com/pytorch/torchtune/blob/main/recipes/configs/llama3_2/1B_full_single_device.yaml, https://github.com/pytorch/torchtune/blob/main/recipes/configs/qwen2/0.5B_full_single_device.yaml.

PPO uses way more memory than many other finetuning algorithms, given limited resources on single device, please consider one configuration using 1B model for PPO.

In terms of default device of CUDA, I can test it on CUDA and update this PR.
Many thanks again. @SalmanMohammadi @ebsmothers

[ebsmothers]

@songhappy thanks, that makes sense. Given it's already 1B, can we use the Llama 3.2 1B model instead? I think this should give better results anyways

[SalmanMohammadi]

@songhappy thanks, that makes sense. Given it's already 1B, can we use the Llama 3.2 1B model instead? I think this should give better results anyways

We don't have a classifier builder for 3.2. We can land #2356 soonish to enable this.