Skip to content

[xpu][float8] Enable float8 fsdp2 tests for Intel XPU #4510

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
arlesniak wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

[xpu][float8] Enable float8 fsdp2 tests for Intel XPU #4510

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9484292
[xpu][float8] Enable float8 fsdp2 for Intel XPU
arlesniak Jun 18, 2026
3ad5700
[xpu][float8] enabled folder test
arlesniak Jun 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 37 additions & 24 deletions test/float8/test_fsdp2/test_fsdp2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
import torch.nn as nn
from torch.distributed._composable.fsdp import MixedPrecisionPolicy, fully_shard
from torch.distributed._tensor import DTensor, init_device_mesh
from torch.testing._internal.common_cuda import TEST_CUDA
from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
from torch.testing._internal.common_fsdp import (
MLP,
Expand Down Expand Up @@ -47,7 +46,10 @@
)
from torchao.utils import is_MI300, is_MI350, is_sm_at_least_89

if not (is_sm_at_least_89() or is_MI300() or is_MI350()):
if not torch.accelerator.is_available():
pytest.skip("GPU not available", allow_module_level=True)

if torch.cuda.is_available() and not (is_sm_at_least_89() or is_MI300() or is_MI350()):
pytest.skip(
"Requires FP8-capable GPU (CUDA SM89+, MI300, or MI350)",
allow_module_level=True,
Expand All @@ -63,13 +65,15 @@ def broadcast_module(self, module: nn.Module) -> None:

def init_single_module(self) -> nn.Module:
torch.manual_seed(42)
module = nn.Linear(16, 16, device="cuda")
device = torch.accelerator.current_accelerator()
module = nn.Linear(16, 16, device=device)
self.broadcast_module(module)
return module

def init_multi_module(self) -> nn.Module:
torch.manual_seed(42)
module = nn.Sequential(*[MLP(16, device="cuda") for _ in range(3)])
device = torch.accelerator.current_accelerator()
module = nn.Sequential(*[MLP(16, device=device) for _ in range(3)])
self.broadcast_module(module)
return module

Expand All @@ -88,7 +92,8 @@ def init_transformer(
weight_tying=weight_tying,
vocab_size=32,
)
module = Transformer(args).cuda()
device = torch.accelerator.current_accelerator()
module = Transformer(args).to(device)
if dtype is not None:
module = module.to(dtype=dtype)

Expand All @@ -103,15 +108,16 @@ def init_transformer(

def get_local_inp(self, dtype: torch.dtype = torch.float32):
torch.manual_seed(42)
global_inp = torch.randn((16 * self.world_size, 16), device="cuda", dtype=dtype)
device = torch.accelerator.current_accelerator()
global_inp = torch.randn((16 * self.world_size, 16), device=device, dtype=dtype)
dist.broadcast(global_inp, src=0)
return global_inp.view(self.world_size, -1)[self.rank].view(16, 16)


class TestFloat8MultiProcess(FSDPTest, TestFloat8Common):
@property
def world_size(self) -> int:
return min(torch.cuda.device_count(), 2)
return min(torch.accelerator.device_count(), 2)

@skip_if_lt_x_gpu(2)
def test_transformer_parity(self):
Expand Down Expand Up @@ -178,8 +184,9 @@ def _test_transformer_parity(
fully_shard(module)
ref_optim = torch.optim.Adam(ref_module.parameters(), lr=1e-2)
optim = torch.optim.Adam(module.parameters(), lr=1e-2, foreach=True)
device = torch.accelerator.current_accelerator()
local_inp = torch.randint(
0, ref_module.tok_embeddings.weight.size(0), (16, 16), device="cuda"
0, ref_module.tok_embeddings.weight.size(0), (16, 16), device=device
)
check_parity_no_mp(
self,
Expand All @@ -201,13 +208,14 @@ def test_transformer_memory(self):

def _test_transformer_memory(self, enable_fsdp_float8_all_gather: bool):
torch.manual_seed(42)
device = torch.accelerator.current_accelerator()
# Pre-run a linear forward (gemm and bias) and backward (gemm) to
# allocate the cuBLAS workspaces before measuring the memory usage
# since the workspace size can differ between hardwares
lin = torch.nn.Linear(768, 768, device="cuda")
inp = torch.randn(1, 768, device="cuda")
lin = torch.nn.Linear(768, 768, device=device)
inp = torch.randn(1, 768, device=device)
lin(inp).sum().backward()
torch.cuda.empty_cache()
torch.accelerator.empty_cache()
base_mem_mb = self._get_peak_active_memory_mb()

vocab_size = 32
Expand Down Expand Up @@ -260,7 +268,7 @@ def _test_transformer_memory(self, enable_fsdp_float8_all_gather: bool):
self.assertLessEqual(curr_mem_mb - base_mem_mb, init_mem_mb)

# Use a small input to minimize activation memory usage
inp = torch.randint(0, vocab_size, (1, 4), device="cuda")
inp = torch.randint(0, vocab_size, (1, 4), device=device)

# Forward:
loss = model(inp)
Expand Down Expand Up @@ -313,11 +321,11 @@ def _test_transformer_memory(self, enable_fsdp_float8_all_gather: bool):
self.assertLessEqual(mem_mb, expected_mem_mb + base_mem_mb)

def _get_peak_active_memory_mb(self) -> int:
mem_stats = torch.cuda.memory_stats()
mem_stats = torch.accelerator.memory_stats()
return round(mem_stats["active_bytes.all.peak"] / 1e6)

def _get_curr_active_memory_mb(self) -> int:
mem_stats = torch.cuda.memory_stats()
mem_stats = torch.accelerator.memory_stats()
return round(mem_stats["active_bytes.all.current"] / 1e6)


Expand All @@ -329,16 +337,17 @@ def world_size(self) -> int:
def test_amax_allreduce_device_mesh(self):
dp_size = 2
pp_size = self.world_size // dp_size
device = torch.accelerator.current_accelerator()
global_mesh = init_device_mesh(
"cuda", (pp_size, dp_size), mesh_dim_names=("pp", "dp")
str(device), (pp_size, dp_size), mesh_dim_names=("pp", "dp")
)
dp_mesh = global_mesh["dp"]

if self.rank in [0, 1]:
# rank 0 and 1 are the 1st stage in the pipeline
# rank 2 and 4 are doing nothing but waiting for the 1st stage
torch.manual_seed(42 + self.rank)
hp_tensor = torch.randn(768, 32, device="cuda")
hp_tensor = torch.randn(768, 32, device=device)
hp_tensor_to_float8_dynamic(
hp_tensor,
e4m3_dtype,
Expand All @@ -356,7 +365,7 @@ class TestFloat8MultiThread(FSDPTestMultiThread, TestFloat8Common):
def world_size(self) -> int:
return 2

@unittest.skipIf(not TEST_CUDA, "no cuda")
@unittest.skipIf(not torch.accelerator.is_available(), "GPU not available")
def test_weight_subclass_dynamic(self):
tensor_cls = WeightWithDynamicFloat8CastTensor
# Check for a single FSDP paramter group
Expand Down Expand Up @@ -393,7 +402,7 @@ def test_weight_subclass_dynamic(self):
if "weight" in param_name:
self.assertIsInstance(param.to_local(), tensor_cls)

@unittest.skipIf(not TEST_CUDA, "no cuda")
@unittest.skipIf(not torch.accelerator.is_available(), "GPU not available")
def test_fp8_fp32_all_gather_dynamic_comm_size(self):
"""
Tests that fp8 all-gather with dynamic scaling communicates the
Expand Down Expand Up @@ -476,7 +485,7 @@ def get_expected_all_gather_size(module: nn.Module):
[s for s in expected_all_gather_sizes for _ in range(self.world_size)],
)

@unittest.skipIf(not TEST_CUDA, "no cuda")
@unittest.skipIf(not torch.accelerator.is_available(), "GPU not available")
def test_fp32_fp8_single_module_parity(self):
"""
Tests numeric parity for fp32 parameters with fp8 computation with a
Expand All @@ -503,7 +512,7 @@ def test_fp32_fp8_single_module_parity(self):
ref_module,
config=float8_linear_config1,
)
ref_module = ref_module.cuda()
ref_module = ref_module.to(torch.accelerator.current_accelerator())
module = convert_to_float8_training(
module_fp32,
config=float8_linear_config2,
Expand All @@ -522,7 +531,7 @@ def test_fp32_fp8_single_module_parity(self):
config=float8_linear_config2,
)

@unittest.skipIf(not TEST_CUDA, "no cuda")
@unittest.skipIf(not torch.accelerator.is_available(), "GPU not available")
def test_fp32_fp8_multi_module_parity(self):
"""
Tests numeric parity for fp32 parameters with fp8 computation with
Expand All @@ -541,7 +550,9 @@ def test_fp32_fp8_multi_module_parity(self):
enable_fsdp_float8_all_gather=enable_fsdp_float8_all_gather,
cast_config_weight=CastConfig(scaling_type=scaling_type_weight),
)
module = self.init_multi_module().cuda()
module = self.init_multi_module().to(
torch.accelerator.current_accelerator()
)
ref_module = copy.deepcopy(module)
ref_module = convert_to_float8_training(
ref_module,
Expand All @@ -567,7 +578,7 @@ def test_fp32_fp8_multi_module_parity(self):
config=float8_linear_config2,
)

@unittest.skipIf(not TEST_CUDA, "no cuda")
@unittest.skipIf(not torch.accelerator.is_available(), "GPU not available")
def test_bf16_mp_fp8_dynamic_multi_parity(self):
"""
Tests numeric parity for fp32 parameters with FSDP's bf16 mixed
Expand All @@ -584,7 +595,9 @@ def test_bf16_mp_fp8_dynamic_multi_parity(self):
ref_module_bf16,
config=float8_config,
)
ref_module_fp32 = copy.deepcopy(module).cuda()
ref_module_fp32 = copy.deepcopy(module).to(
torch.accelerator.current_accelerator()
)
module = convert_to_float8_training(module, config=float8_config)
mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16)
for mlp in module:
Expand Down
7 changes: 4 additions & 3 deletions test/float8/test_fsdp2_tp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,17 +32,18 @@

def setup_distributed():
world_size = int(os.environ.get("WORLD_SIZE", -1))
device = str(torch.accelerator.current_accelerator())

# https://pytorch.org/tutorials/recipes/distributed_device_mesh.html
device_mesh = init_device_mesh(
"cuda",
device,
(world_size // 2, 2),
mesh_dim_names=("dp", "tp"),
)
# seed must be the same in all processes
torch.manual_seed(1)
local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
torch.get_device_module(device).set_device(local_rank)
return device_mesh


Expand Down Expand Up @@ -84,7 +85,7 @@ def _test_fp8_mlp_tensor_parallelism_base(

tp_out = tp_model(x_fp32_tp_input)
tp_out.sum().backward()
torch.cuda.synchronize()
torch.get_device_module(device).synchronize()

# TODO(future PR): test numerics, and add more cases

Expand Down
Loading