From 15a30b7d7a982f58a4324afc58476d546a72203d Mon Sep 17 00:00:00 2001 From: Dmitry Rogozhkin Date: Tue, 8 Apr 2025 11:35:57 -0700 Subject: [PATCH] feat: support cpu and xpu devices in llama4 CPU support tried on Intel Xeon. XPU support tried on Intel Data Center GPU Max series. Note that a fix in fairscale for https://github.com/facebookresearch/fairscale/issues/1195 is required to make XPU working (CPU not affected). I believe that this fix might also be needed for CUDA as well since issue seems to be device agnostic and I see it with the simplified reproducer on NVidia A10. Verified on the platforms named above with LLama4 sample completion and chat completion scripts. Requires: https://github.com/facebookresearch/fairscale/pull/1196 Signed-off-by: Dmitry Rogozhkin --- models/llama4/generation.py | 54 +++++++++++++++++++----- models/llama4/model.py | 4 +- models/llama4/scripts/chat_completion.py | 14 ++++++ models/llama4/scripts/completion.py | 14 ++++++ 4 files changed, 74 insertions(+), 12 deletions(-) diff --git a/models/llama4/generation.py b/models/llama4/generation.py index 67dfad460..692f4fa45 100644 --- a/models/llama4/generation.py +++ b/models/llama4/generation.py @@ -11,6 +11,7 @@ import os import sys import time +from packaging import version from pathlib import Path from typing import Callable, Generator, List, Optional @@ -33,6 +34,12 @@ torch.serialization.add_safe_globals([io.BytesIO, codecs.encode]) +def is_xccl_available(): + if version.parse(torch.__version__).release >= version.parse("2.7").release: + return torch.distributed.distributed_c10d.is_xccl_available() + return False + + class Llama4: @staticmethod def build( @@ -42,9 +49,24 @@ def build( world_size: Optional[int] = None, quantization_mode: Optional[QuantizationMode] = None, seed: int = 1, + device: str = "cuda", ): + device = torch.device(device) + if ( + device.type == "cuda" + and not torch.cuda.is_available() + or device.type == "xpu" + and not torch.xpu.is_available() + ): + raise RuntimeError(f"PyTorch backend for {device.type} device type is not available") + if not torch.distributed.is_initialized(): - torch.distributed.init_process_group("nccl") + if device.type == "cuda": + torch.distributed.init_process_group("nccl") + elif device.type == "xpu" and is_xccl_available(): + torch.distributed.init_process_group("xccl") + else: + torch.distributed.init_process_group("gloo") if not model_parallel_is_initialized(): if world_size is None: @@ -52,7 +74,10 @@ def build( initialize_model_parallel(world_size) local_rank = int(os.environ.get("LOCAL_RANK", 0)) - torch.cuda.set_device(local_rank) + if device.type == "cuda": + torch.cuda.set_device(local_rank) + elif device.type == "xpu": + torch.xpu.set_device(local_rank) torch.manual_seed(seed) @@ -96,15 +121,24 @@ def build( print("Done...") model = convert_to_quantized_model(model, ckpt_dir, quantization_mode) else: - if torch.cuda.is_bf16_supported(): - torch.set_default_tensor_type(torch.cuda.BFloat16Tensor) - else: - torch.set_default_tensor_type(torch.cuda.HalfTensor) + print(f"Setting default device to {device}") + torch.set_default_device(device) + if device.type == "cuda": + if torch.cuda.is_bf16_supported(): + torch.set_default_dtype(torch.bfloat16) + else: + torch.set_default_dtype(torch.half) + elif device.type == "xpu": + if torch.xpu.is_bf16_supported(): + torch.set_default_dtype(torch.bfloat16) + else: + torch.set_default_dtype(torch.half) model = Transformer(model_args) print("Loading state dict...") model.load_state_dict(state_dict, strict=False) print("Done...") + model.to(device) print(f"Loaded in {time.time() - start_time:.2f} seconds") return Llama4(model, tokenizer, model_args) @@ -152,13 +186,13 @@ def generate( total_len = min(max_gen_len + max_prompt_len, params.max_seq_len) pad_id = self.tokenizer.pad_id - tokens = torch.full((bsz, total_len), pad_id, dtype=torch.long, device="cuda") + tokens = torch.full((bsz, total_len), pad_id, dtype=torch.long) for k, t in enumerate(prompt_tokens): - tokens[k, : len(t)] = torch.tensor(t, dtype=torch.long, device="cuda") + tokens[k, : len(t)] = torch.tensor(t, dtype=torch.long) if logprobs: token_logprobs = torch.zeros_like(tokens, dtype=torch.float) - eos_reached = torch.tensor([False] * bsz, device="cuda") + eos_reached = torch.tensor([False] * bsz) input_text_mask = tokens != pad_id if echo: @@ -178,7 +212,7 @@ def generate( ) yield results - stop_tokens = torch.tensor(self.tokenizer.stop_tokens, device="cuda") + stop_tokens = torch.tensor(self.tokenizer.stop_tokens) prev_pos = 0 for cur_pos in range(min_prompt_len, total_len): diff --git a/models/llama4/model.py b/models/llama4/model.py index 75f281d5c..272164142 100644 --- a/models/llama4/model.py +++ b/models/llama4/model.py @@ -161,7 +161,7 @@ def __init__( self.n_local_kv_heads, self.head_dim, ) - ).cuda() + ) self.cache_v = torch.zeros( ( args.max_batch_size, @@ -169,7 +169,7 @@ def __init__( self.n_local_kv_heads, self.head_dim, ) - ).cuda() + ) self.norm_eps = args.norm_eps self._register_load_state_dict_pre_hook(self.load_hook) diff --git a/models/llama4/scripts/chat_completion.py b/models/llama4/scripts/chat_completion.py index 458c23830..d14d5dca1 100644 --- a/models/llama4/scripts/chat_completion.py +++ b/models/llama4/scripts/chat_completion.py @@ -18,9 +18,22 @@ from models.datatypes import RawMediaItem, RawMessage, RawTextItem, StopReason from models.llama4.generation import Llama4 +import os +import torch + THIS_DIR = Path(__file__).parent +def get_device(): + if "DEVICE" in os.environ: + return os.environ["DEVICE"] + if torch.cuda.is_available(): + return "cuda" + elif torch.xpu.is_available(): + return "xpu" + return "cpu" + + def run_main( checkpoint_dir: str, world_size: int = 1, @@ -36,6 +49,7 @@ def run_main( max_batch_size=max_batch_size, world_size=world_size, quantization_mode=quantization_mode, + device=get_device(), ) dialogs = [ diff --git a/models/llama4/scripts/completion.py b/models/llama4/scripts/completion.py index ada3669b8..96a9f32a1 100644 --- a/models/llama4/scripts/completion.py +++ b/models/llama4/scripts/completion.py @@ -18,9 +18,22 @@ from models.datatypes import RawMediaItem from models.llama4.generation import Llama4 +import os +import torch + THIS_DIR = Path(__file__).parent +def get_device(): + if "DEVICE" in os.environ: + return os.environ["DEVICE"] + if torch.cuda.is_available(): + return "cuda" + elif torch.xpu.is_available(): + return "xpu" + return "cpu" + + def run_main( checkpoint_dir: str, world_size: int = 1, @@ -36,6 +49,7 @@ def run_main( max_batch_size=max_batch_size, world_size=world_size, quantization_mode=quantization_mode, + device=get_device(), ) with open(THIS_DIR / "../../resources/dog.jpg", "rb") as f: