profiling ops on xpu

[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-2875

Changelog

What are the changes made in this PR?
Added 'xpu' in _profiler.py and modify cuda related memory profiler to cuda only in driver scripts in receipts directory.
*

Test plan

Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.

• manually run any new or modified recipes with sufficient proof of correctness
  Steps:

1. download Llama-3.2-3B-Instruct model
2. modify recipes/configs/llama3_2/3B_full_single_device.yaml, change "device: xpu", "profiler.enabled:True"
3. tune run full_finetune_single_device --config recipes/configs/llama3_2/3B_full_single_device.yaml
4. see profiling results under /tmp/full-llama3.2-finetune/profiling_results

{
  "schemaVersion": 1,
  "deviceProperties": [
  ],
  "with_flops": 1,
  "record_shapes": 1,
  "profile_memory": 1,
  "with_stack": 1,
  "trace_id": "1DB69D6280304432870B411620032DC3",
  "traceEvents": [
  {
    "ph": "X", "cat": "cpu_op", "name": "aten::conv2d", "pid": 1341316, "tid": 1
341316,
    "ts": 731713074715.025, "dur": 95265.387,
    "args": {
      "External id": 1,"Record function id": 0, "Concrete Inputs": ["", "", "",
"[2, 2]", "[3, 3]", "[1, 1]", "1"], "Input type": ["float", "float", "", "Scalar
List", "ScalarList", "ScalarList", "Scalar"], "Input Strides": [[150528, 50176,
224, 1], [147, 49, 7, 1], [], [], [], [], []], "Input Dims": [[32, 3, 224, 224],
 [64, 3, 7, 7], [], [], [], [], []], "Ev Idx": 0
    }
  },

UX

If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example

• I did not change any public API

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[songhappy]

@SalmanMohammadi Could you please review and approve it?

[felipemello1]

hey @songhappy , thanks for the PR!

Just two questions before i approve it:

1. i saw that you added and self._device.type == "cuda" to some recipes, but not all. Is it because the other recipes already have it, or did we forget some?

2. In your testing, i dont think that you set profiler.profile_memory=True. Do you think its worth checking this option to make sure its working for xpu?

[songhappy]

@felipemello1 thanks a lot for reviewing it.

1. Yes, all of them should have and self._device.type == "cuda", I forgot one, and added it now.
2. modified testcases to cover profiler.profile_memory=True on xpu too.
3. I have run recipe while setting profiler.profile_memory=True in the configeration file and copied some profiling logging when add this PR, it is not necessary to set profiler.profile_memory=True as default.

[felipemello1]

Awesome, thanks @songhappy , do you mind just adding to https://github.com/pytorch/torchtune/blob/main/recipes/ppo_full_finetune_single_device.py?

I will merge after it

[songhappy]

@felipemello1 added in ppo. One question, have you tried smaller models other than 7b to run https://github.com/pytorch/torchtune/blob/main/recipes/ppo_full_finetune_single_device.py? like 3b llama or similar models. I want to try smaller models of PPO. Please guide