Adding Intel XPU backend support using device-agnostic way #1280

shiyan1121 · 2024-08-07T06:54:31Z

Context

This PR enhances the Torchtune framework by introducing support for the Intel XPU backend.

Changelog

To avoid the redundant hard-code for XPU logic and simplify the device abstraction, we follows the implementation in PyTorch FSDP:
Making fsdp device-agnostic for custom-backend which implement cuda-semantics #99024

In the torchtune/utils/_device.py file, we introduce the _DeviceHandle class. Users can initialize this class with their device by using device_handle = _DeviceHandle.from_device(my_device). Since XPU also follows the CUDA semantics and provides CUDA-like APIs, e.g. torch.xpu.empty_cache(), most torch.cuda. semantics could be replaced by device_handle., which are available both for CUDA and XPU, and even for other CUDA-like devices.

For some logics that the device_handle strategy cannot be implemented, we resort to using XPU related judgement to handle different scenarios.

This reverts commit 8203a3b.

pytorch-bot · 2024-08-07T06:54:34Z

🔗 Helpful Links

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

📄 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.

facebook-github-bot · 2024-08-07T06:54:37Z

Hi @HJG971121!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

ebsmothers · 2024-08-09T05:24:28Z

Thanks @HJG971121 for the PR. Overall I think the _DeviceHandle abstraction is pretty nice.

One very basic question is around the usage of intel_extension_for_pytorch and oneccl_bindings_for_pytorch, are these just defining all the necessary bindings to the XPU backend? And if so, do they need to be imported globally, or can we put them in a utility somewhere?

And a second question: does torchtune run on XPU with the changes in this PR? It'd be nice to add some sample runs confirming things run and have reasonable performance.

Add global xpu import in utils.

shiyan1121 · 2024-08-13T09:25:46Z

Hi, @ebsmothers, thank you for your response.

Regarding your first question, both intel_extension_for_pytorch and oneccl_bindings_for_pytorch are essential for our XPU backend. I have incorporated them into utils for global import to ensure they are readily accessible throughout the project.

For your second question, we have successfully tested the Torchtune functionality with the llama2-7b and llama3-8b models on our XPU. Attached is the log for fine-tuning the llama2-7b model. Could you advise if it would be appropriate to include our recipe in the current PR?
llama2-7b_single_device_lora_bs_4_seqlen_2048_test_data_pack.log

shiyan1121 · 2024-08-21T03:09:49Z

Here is our example of command line for llama2-7b full finetuning on single device, where the device is specified as XPU:

tune run full_finetune_single_device --config llama2/7B_full_low_memory device=xpu

facebook-github-bot · 2024-08-21T04:06:55Z

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

RdoubleA

Like the idea of the DeviceHandle a lot. Had a few questions around its usage.

RdoubleA · 2024-08-21T17:34:58Z

recipes/full_finetune_distributed.py

            param_init_fn=(
                lambda module: module.to_empty(
-                    device=torch.device("cuda"), recurse=False
+                    device=torch.device("xpu") if torch.xpu.is_available() else torch.device("cuda"), recurse=False


If device_handle is serving as an intermediary wrapper to handle all device related APIs agnostic of whether that is CPU, CUDA, or XPU, shouldn't we work with device_handle directly here instead of torch.device? My understanding is that we only need to find the appropriate device one time in the recipe (via get_device), then call get_device_handle, and use the device_handle for all these operations. Or am I misunderstanding?

Thank you for your comment!
Your understanding is correct，and it's my mistake that I did't use unified device handle in this instance.

Additionally, I suggest that code here could be rewritten as
module.to_empty(device = self._device.type, recurse=False)
to avoid redundant device verification, given that the device attribute is already specified.

RdoubleA · 2024-08-21T17:35:08Z

recipes/full_finetune_distributed.py

        )
-
-    init_process_group(backend="gloo" if cfg.device == "cpu" else "nccl")
+    if torch.xpu.is_available():


same comment here

Thank you for pointing that out!

Indeed, there are scenarios where the device_handle might not be applicable, for example backend needs to be determined based on the device type.

And I think using cfg.device as the condition variable may be a simpler way. I will rewrite the code, thank you.

RdoubleA · 2024-08-21T17:37:01Z

recipes/lora_finetune_single_device.py

-                                self._device.type == "cuda"
-                                and self._log_peak_memory_stats
-                            ):
+                            if self._log_peak_memory_stats:


do we want to guard against CPU here? get_memory_stats will throw an error but maybe we shouldn't log at all if on CPU. (actually, none of our recipes support CPU training, so maybe we should just let it error out, @ebsmothers)

RdoubleA · 2024-08-21T17:38:05Z

torchtune/utils/_device.py

+    """
+    This is a simple abstraction for computing devices,
+    which enables custom backends that implement CUDA-like
+    semantics.


Could you provide a few more details here on how this should be used? MY understanding is that this is used in place of device. A few examples would be helpful.

This device_handleis designed to eliminate redundant code when implementing CUDA-like APIs for your customized device. For instance, consider the XPU. To empty the cache, a simplistic approach involves using conditional statements:

torch.xpu.empty_cache() if device == 'xpu' else torch.cuda.empty_cache()

This method, however, lacks efficiency and flexibility, particularly as the number of GPU devices increases in the future. Therefore, we adopt the device_handle strategy to integrate both CUDA APIs and the CUDA-like APIs of the customized device.

To utilize this strategy, you should first define your device_handle based on your device:

device_handle = get_device_handle (device) # device could be ‘cuda’, ‘xpu’ or your customed device

Then, the code for emptying the cache can be rewritten as:

device_handle.empty_cache()

In this setup, device_handle replaces direct calls to torch.cuda or torch.xpu, streamlining the process across different device types. Commonly used CUDA APIs, such as torch.cuda.memory_allocated() and torch.cuda.synchronize(), can be utilized in the same manner.

RdoubleA · 2024-08-21T17:39:25Z

torchtune/utils/_device.py

+            return cast(_DeviceHandle, torch.cuda)
+        return cls(device)
+
+    def __getattr__(self, __name: str) -> Any:


a quick docstring here would be helpful, something along the lines of "retrieve attribute from correct backend based on registered device"

Thank you for your suggestion! I will add docstring to it.

RdoubleA · 2024-08-21T17:39:49Z

torchtune/utils/_device.py

-    torch.cuda.set_device(device)
+    device_handle = get_device_handle(device)
+    if device.index >= device_handle.device_count():
+            raise RuntimeError(


nit: extra tab here

zhuhong61 · 2024-08-29T06:50:00Z

Convert to Draft. We are working on preparing the RFC for the discussion about the Intel PyTorch distributed design. We will push forward the process after reaching an alignment :)

Ankur-singh · 2025-01-06T06:18:29Z

Convert to Draft. We are working on preparing the RFC for the discussion about the Intel PyTorch distributed design. We will push forward the process after reaching an alignment :)

Hi @zhuhong61 any updates on this?

songhappy · 2025-01-22T22:43:52Z

@shiyan1121 Could you close this one please? New PRs have been added according to new code to support xpu. see #1953, #1826, #2249

zhuhong61 and others added 7 commits June 18, 2024 04:41

Support XPU device for Llama2 fine-tuning

1225ebd

XPU backend using device handle.

31c8d73

delete recipes for inner developers

6ab4263

modify some details

13fc377

Modify some details.

8203a3b

Revert "Modify some details."

95baa09

This reverts commit 8203a3b.

Modify some details.

fa6ce16

joecummings requested review from ebsmothers and pbontrager August 9, 2024 01:42

shiyan1121 and others added 2 commits August 13, 2024 01:50

Add global xpu import in utils.

7d7d322

Merge pull request #2 from zhuhong61/sy/xpu_device_handle

cff49d3

Add global xpu import in utils.

felipemello1 mentioned this pull request Aug 18, 2024

Improve training UX with TPS, GPU peak mem %, cudaMalloc retries, and do it in color for pizzaz #1360

Open

12 tasks

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 21, 2024

RdoubleA mentioned this pull request Aug 21, 2024

Support for Intel GPUs #999

Closed

RdoubleA reviewed Aug 21, 2024

View reviewed changes

RdoubleA mentioned this pull request Aug 21, 2024

how to support huawei Ascend npu device #1006

Closed

shiyan1121 marked this pull request as draft August 29, 2024 06:12

RdoubleA mentioned this pull request Nov 16, 2024

Can we get XPU Nightlies? #2005

Closed

Jack-Khuu mentioned this pull request Jan 15, 2025

Add Intel XPU device support to generate and serve pytorch/torchchat#1361

Merged

SalmanMohammadi closed this Mar 28, 2025

Adding Intel XPU backend support using device-agnostic way #1280

Adding Intel XPU backend support using device-agnostic way #1280

Uh oh!

Conversation

shiyan1121 commented Aug 7, 2024

Context

Changelog

Uh oh!

pytorch-bot bot commented Aug 7, 2024

🔗 Helpful Links

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

Uh oh!

facebook-github-bot commented Aug 7, 2024

Action Required

Process

Uh oh!

ebsmothers commented Aug 9, 2024

Uh oh!

shiyan1121 commented Aug 13, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shiyan1121 commented Aug 21, 2024

Uh oh!

facebook-github-bot commented Aug 21, 2024

Uh oh!

RdoubleA left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhuhong61 commented Aug 29, 2024

Uh oh!

Ankur-singh commented Jan 6, 2025

Uh oh!

songhappy commented Jan 22, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

shiyan1121 commented Aug 13, 2024 •

edited

Loading