mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-23 09:09:01 +03:00
ee849197ce
* feat: qwen image 21 tp plan * test: vallidated on neuron * fix: lru_cache_unless_export leak * fix: use complex RoPE on TPU for Qwen Image 2.1 TorchTPU reports TPU tensors as "tpu" and supports complex dtypes, so only Neuron needs the angle-based RoPE path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * refactor: inline the Neuron device check for Qwen Image 2.1 RoPE * review: cache Qwen Image 2.1 RoPE device freqs with lru_cache * review: drop redundant torch.device * fix: use torch.accelerator.device_count() when resolving the parallel device * refactor: drop the Neuron branch in Qwen Image 2.1 block ids * restore: neuron repeat_interleave lowering error fix * review: add TensorParallelNeuronTesterMixin for Neuron TP tests * review: revert Qwen Image 2.1 RoPE freqs cache to a per-instance dict * review: drop the Neuron RoPE path in Qwen Image 2.1 Neuron now runs complex tensors natively, so Qwen Image 2.1 uses the default complex RoPE on every device. This removes `apply_rotary_emb_qwen_neuron`, `ROPE_PER_DEVICE` and the per-device freqs cache, and keeps `self.freqs = [freq.to(device) ...]` as on main. --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
881 lines
34 KiB
Python
881 lines
34 KiB
Python
# coding=utf-8
|
|
# Copyright 2026 HuggingFace Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import os
|
|
import socket
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.multiprocessing as mp
|
|
|
|
from diffusers.models._modeling_parallel import ContextParallelConfig, TensorParallelConfig
|
|
from diffusers.models.attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
|
|
from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME
|
|
|
|
from ...testing_utils import (
|
|
is_attention,
|
|
is_context_parallel,
|
|
is_kernels_available,
|
|
is_tensor_parallel,
|
|
require_torch_multi_accelerator,
|
|
require_torch_neuron,
|
|
require_torch_tpu,
|
|
torch_device,
|
|
)
|
|
from .common import calculate_expected_num_shards, compute_module_persistent_sizes
|
|
from .utils import _maybe_cast_to_bf16
|
|
|
|
|
|
# Device configuration mapping
|
|
DEVICE_CONFIG = {
|
|
"cuda": {"backend": "nccl", "module": torch.cuda},
|
|
"xpu": {"backend": "xccl", "module": torch.xpu},
|
|
"tpu": {"backend": "tpu_dist", "module": None},
|
|
"neuron": {"backend": "neuron", "module": getattr(torch, "neuron", None)},
|
|
}
|
|
|
|
|
|
def _find_free_port():
|
|
"""Find a free port on localhost."""
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
s.bind(("", 0))
|
|
s.listen(1)
|
|
port = s.getsockname()[1]
|
|
return port
|
|
|
|
|
|
def _context_parallel_worker(
|
|
rank,
|
|
world_size,
|
|
master_port,
|
|
model_class,
|
|
init_dict,
|
|
cp_dict,
|
|
inputs_dict,
|
|
return_dict,
|
|
attention_backend=None,
|
|
state_dict=None,
|
|
):
|
|
"""Worker function for context parallel testing."""
|
|
try:
|
|
# Set up distributed environment
|
|
os.environ["MASTER_ADDR"] = "localhost"
|
|
os.environ["MASTER_PORT"] = str(master_port)
|
|
os.environ["RANK"] = str(rank)
|
|
os.environ["WORLD_SIZE"] = str(world_size)
|
|
|
|
# Get device configuration
|
|
device_config = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])
|
|
backend = device_config["backend"]
|
|
device_module = device_config["module"]
|
|
|
|
# Initialize process group
|
|
dist.init_process_group(backend=backend, rank=rank, world_size=world_size)
|
|
|
|
# Set device for this process
|
|
device_module.set_device(rank)
|
|
device = torch.device(f"{torch_device}:{rank}")
|
|
|
|
# Create model
|
|
model = model_class(**init_dict)
|
|
if state_dict is not None:
|
|
model.load_state_dict(state_dict)
|
|
model.to(device)
|
|
model.eval()
|
|
|
|
# Cast as needed.
|
|
model, inputs_dict = _maybe_cast_to_bf16(attention_backend, model, inputs_dict)
|
|
|
|
# Move inputs to device
|
|
inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
|
|
|
|
# Enable attention backend
|
|
if attention_backend:
|
|
model.set_attention_backend(attention_backend)
|
|
|
|
# Enable context parallelism
|
|
cp_config = ContextParallelConfig(**cp_dict)
|
|
model.enable_parallelism(config=cp_config)
|
|
|
|
# Run forward pass
|
|
with torch.no_grad():
|
|
output = model(**inputs_on_device, return_dict=False)[0]
|
|
|
|
# Only rank 0 reports results
|
|
if rank == 0:
|
|
return_dict["status"] = "success"
|
|
return_dict["output_shape"] = list(output.shape)
|
|
if state_dict is not None:
|
|
# Serialise via nested list so the manager dict can transport it across processes.
|
|
return_dict["output"] = output.cpu().tolist()
|
|
|
|
except Exception as e:
|
|
if rank == 0:
|
|
return_dict["status"] = "error"
|
|
return_dict["error"] = str(e)
|
|
finally:
|
|
if dist.is_initialized():
|
|
dist.destroy_process_group()
|
|
|
|
|
|
def _context_parallel_backward_worker(
|
|
rank, world_size, master_port, model_class, init_dict, cp_dict, inputs_dict, return_dict
|
|
):
|
|
"""Worker function for context parallel backward pass testing."""
|
|
try:
|
|
# Set up distributed environment
|
|
os.environ["MASTER_ADDR"] = "localhost"
|
|
os.environ["MASTER_PORT"] = str(master_port)
|
|
os.environ["RANK"] = str(rank)
|
|
os.environ["WORLD_SIZE"] = str(world_size)
|
|
|
|
# Get device configuration
|
|
device_config = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])
|
|
backend = device_config["backend"]
|
|
device_module = device_config["module"]
|
|
|
|
# Initialize process group
|
|
dist.init_process_group(backend=backend, rank=rank, world_size=world_size)
|
|
|
|
# Set device for this process
|
|
device_module.set_device(rank)
|
|
device = torch.device(f"{torch_device}:{rank}")
|
|
|
|
# Create model in training mode
|
|
model = model_class(**init_dict)
|
|
model.to(device)
|
|
model.train()
|
|
|
|
# Move inputs to device
|
|
inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
|
|
|
|
# Enable context parallelism
|
|
cp_config = ContextParallelConfig(**cp_dict)
|
|
model.enable_parallelism(config=cp_config)
|
|
|
|
# Run forward and backward pass
|
|
output = model(**inputs_on_device, return_dict=False)[0]
|
|
loss = output.sum()
|
|
loss.backward()
|
|
|
|
# Check that backward actually produced at least one valid gradient
|
|
grads = [p.grad for p in model.parameters() if p.requires_grad and p.grad is not None]
|
|
has_valid_grads = len(grads) > 0 and all(torch.isfinite(g).all() for g in grads)
|
|
|
|
# Only rank 0 reports results
|
|
if rank == 0:
|
|
return_dict["status"] = "success"
|
|
return_dict["has_valid_grads"] = bool(has_valid_grads)
|
|
|
|
except Exception as e:
|
|
if rank == 0:
|
|
return_dict["status"] = "error"
|
|
return_dict["error"] = str(e)
|
|
finally:
|
|
if dist.is_initialized():
|
|
dist.destroy_process_group()
|
|
|
|
|
|
def _custom_mesh_worker(
|
|
rank,
|
|
world_size,
|
|
master_port,
|
|
model_class,
|
|
init_dict,
|
|
cp_dict,
|
|
mesh_shape,
|
|
mesh_dim_names,
|
|
inputs_dict,
|
|
return_dict,
|
|
):
|
|
"""Worker function for context parallel testing with a user-provided custom DeviceMesh."""
|
|
try:
|
|
os.environ["MASTER_ADDR"] = "localhost"
|
|
os.environ["MASTER_PORT"] = str(master_port)
|
|
os.environ["RANK"] = str(rank)
|
|
os.environ["WORLD_SIZE"] = str(world_size)
|
|
|
|
# Get device configuration
|
|
device_config = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])
|
|
backend = device_config["backend"]
|
|
device_module = device_config["module"]
|
|
|
|
dist.init_process_group(backend=backend, rank=rank, world_size=world_size)
|
|
|
|
# Set device for this process
|
|
device_module.set_device(rank)
|
|
device = torch.device(f"{torch_device}:{rank}")
|
|
|
|
model = model_class(**init_dict)
|
|
model.to(device)
|
|
model.eval()
|
|
|
|
inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
|
|
|
|
# DeviceMesh must be created after init_process_group, inside each worker process.
|
|
mesh = torch.distributed.device_mesh.init_device_mesh(
|
|
torch_device, mesh_shape=mesh_shape, mesh_dim_names=mesh_dim_names
|
|
)
|
|
cp_config = ContextParallelConfig(**cp_dict, mesh=mesh)
|
|
model.enable_parallelism(config=cp_config)
|
|
|
|
with torch.no_grad():
|
|
output = model(**inputs_on_device, return_dict=False)[0]
|
|
|
|
if rank == 0:
|
|
return_dict["status"] = "success"
|
|
return_dict["output_shape"] = list(output.shape)
|
|
|
|
except Exception as e:
|
|
if rank == 0:
|
|
return_dict["status"] = "error"
|
|
return_dict["error"] = str(e)
|
|
finally:
|
|
if dist.is_initialized():
|
|
dist.destroy_process_group()
|
|
|
|
|
|
def _tensor_parallel_worker(
|
|
rank,
|
|
world_size,
|
|
master_port,
|
|
model_class,
|
|
init_dict,
|
|
inputs_dict,
|
|
return_dict,
|
|
state_dict,
|
|
device_type=torch_device,
|
|
):
|
|
"""Worker function for tensor parallel inference testing.
|
|
|
|
Each rank builds the (identical, `state_dict`-loaded) model, sets up the accelerator device, shards the model
|
|
with `enable_parallelism(config=TensorParallelConfig(tp_degree=world_size))` and runs a forward pass. Rank 0
|
|
reports its output so the caller can compare it against a single-device reference (TP is mathematically equivalent
|
|
to the unsharded model up to floating-point reduction order).
|
|
"""
|
|
try:
|
|
os.environ["MASTER_ADDR"] = "localhost"
|
|
os.environ["MASTER_PORT"] = str(master_port)
|
|
os.environ["RANK"] = str(rank)
|
|
os.environ["LOCAL_RANK"] = str(rank)
|
|
os.environ["WORLD_SIZE"] = str(world_size)
|
|
|
|
device_config = DEVICE_CONFIG.get(device_type, DEVICE_CONFIG["cuda"])
|
|
dist.init_process_group(backend=device_config["backend"], rank=rank, world_size=world_size)
|
|
|
|
if device_type == "tpu":
|
|
# Each spawned process is bound to one chip. Avoid bf16 matmuls to keep the tolerance tight.
|
|
device = torch.device("tpu")
|
|
torch.set_float32_matmul_precision("highest")
|
|
else:
|
|
device_config["module"].set_device(rank)
|
|
device = torch.device(f"{device_type}:{rank}")
|
|
|
|
model = model_class(**init_dict)
|
|
model.load_state_dict(state_dict)
|
|
model.to(device)
|
|
model.eval()
|
|
|
|
inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
|
|
|
|
# Shard the model across all ranks; the device mesh is built from `tp_degree` on the active accelerator.
|
|
model.enable_parallelism(config=TensorParallelConfig(tp_degree=world_size))
|
|
|
|
with torch.no_grad():
|
|
output = model(**inputs_on_device, return_dict=False)[0]
|
|
|
|
if rank == 0:
|
|
return_dict["status"] = "success"
|
|
return_dict["output_shape"] = list(output.shape)
|
|
# Serialise via nested list so the manager dict can transport it across processes.
|
|
return_dict["output"] = output.float().cpu().tolist()
|
|
|
|
except Exception as e:
|
|
if rank == 0:
|
|
return_dict["status"] = "error"
|
|
return_dict["error"] = str(e)
|
|
finally:
|
|
if dist.is_initialized():
|
|
dist.destroy_process_group()
|
|
|
|
|
|
def _tensor_parallel_from_pretrained_worker(
|
|
rank, world_size, master_port, model_class, checkpoint_dir, inputs_dict, return_dict
|
|
):
|
|
"""Worker for `from_pretrained(..., parallel_config=...)`, i.e. sharding while reading the checkpoint.
|
|
|
|
Each rank loads only its own slice of every `_tp_plan` parameter straight into a `DTensor` and runs a forward
|
|
pass. Rank 0 checks that exactly the parameters `_tp_plan` covers were loaded as `DTensor`s, each with the
|
|
placement and local shape its shard spec implies, and reports its output so the caller can check the numerics.
|
|
"""
|
|
try:
|
|
os.environ["MASTER_ADDR"] = "localhost"
|
|
os.environ["MASTER_PORT"] = str(master_port)
|
|
os.environ["RANK"] = str(rank)
|
|
os.environ["WORLD_SIZE"] = str(world_size)
|
|
|
|
device_config = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])
|
|
dist.init_process_group(backend=device_config["backend"], rank=rank, world_size=world_size)
|
|
device_config["module"].set_device(rank)
|
|
|
|
from torch.distributed.tensor import DTensor, Replicate, Shard
|
|
|
|
from diffusers.hooks.tensor_parallel import resolve_tp_shard_specs
|
|
|
|
model = model_class.from_pretrained(
|
|
checkpoint_dir, parallel_config=TensorParallelConfig(tp_degree=world_size)
|
|
).eval()
|
|
|
|
device = torch.device(f"{torch_device}:{rank}")
|
|
inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
|
|
with torch.no_grad():
|
|
output = model(**inputs_on_device, return_dict=False)[0]
|
|
if isinstance(output, DTensor):
|
|
output = output.full_tensor()
|
|
|
|
if rank == 0:
|
|
specs = resolve_tp_shard_specs(model, model_class._tp_plan, world_size)
|
|
state_dict = model.state_dict()
|
|
|
|
# The numerics check alone wouldn't catch a planned parameter left unsharded.
|
|
for name, spec in specs.items():
|
|
param = state_dict[name]
|
|
assert isinstance(param, DTensor), (
|
|
f"'{name}' is covered by `_tp_plan` but was not loaded as a DTensor."
|
|
)
|
|
placement = Replicate() if spec.dim is None else Shard(spec.dim)
|
|
assert param.placements == (placement,), (
|
|
f"'{name}' has placements {param.placements}, not {placement}."
|
|
)
|
|
expected_local_shape = list(param.shape)
|
|
if spec.dim is not None:
|
|
expected_local_shape[spec.dim] //= world_size
|
|
assert list(param.to_local().shape) == expected_local_shape, (
|
|
f"'{name}' has local shape {list(param.to_local().shape)}, not {expected_local_shape}."
|
|
)
|
|
|
|
unplanned = sorted(k for k, v in state_dict.items() if isinstance(v, DTensor) and k not in specs)
|
|
assert not unplanned, f"Parameters not covered by `_tp_plan` were loaded as DTensors: {unplanned}"
|
|
|
|
return_dict["status"] = "success"
|
|
return_dict["output"] = output.float().cpu().tolist()
|
|
|
|
except Exception as e:
|
|
if rank == 0:
|
|
return_dict["status"] = "error"
|
|
return_dict["error"] = f"{type(e).__name__}: {e}"
|
|
finally:
|
|
if dist.is_initialized():
|
|
dist.destroy_process_group()
|
|
|
|
|
|
@is_tensor_parallel
|
|
@require_torch_multi_accelerator
|
|
class TensorParallelTesterMixin:
|
|
def test_tensor_parallel_inference(self, batch_size: int = 1):
|
|
if not torch.distributed.is_available():
|
|
pytest.skip("torch.distributed is not available.")
|
|
|
|
if getattr(self.model_class, "_tp_plan", None) is None:
|
|
pytest.skip("Model does not define a `_tp_plan` for tensor parallel inference.")
|
|
|
|
world_size = 2
|
|
init_dict = self.get_init_dict()
|
|
num_heads = init_dict.get("num_attention_heads")
|
|
if num_heads is not None and num_heads % world_size != 0:
|
|
pytest.skip(f"`num_attention_heads` ({num_heads}) is not divisible by tp_degree ({world_size}).")
|
|
|
|
inputs_dict = self.get_dummy_inputs(batch_size=batch_size)
|
|
|
|
# Single-device reference
|
|
model = self.model_class(**init_dict).eval().to(torch_device)
|
|
state_dict = {k: v.cpu() for k, v in model.state_dict().items()}
|
|
with torch.no_grad():
|
|
ref_output = model(**inputs_dict, return_dict=False)[0].float().cpu()
|
|
|
|
# Move all tensors to CPU for multiprocessing
|
|
inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
|
|
|
|
master_port = _find_free_port()
|
|
manager = mp.Manager()
|
|
return_dict = manager.dict()
|
|
|
|
mp.spawn(
|
|
_tensor_parallel_worker,
|
|
args=(world_size, master_port, self.model_class, init_dict, inputs_dict, return_dict, state_dict),
|
|
nprocs=world_size,
|
|
join=True,
|
|
)
|
|
|
|
assert return_dict.get("status") == "success", (
|
|
f"Tensor parallel inference failed: {return_dict.get('error', 'Unknown error')}"
|
|
)
|
|
|
|
tp_output = torch.tensor(return_dict["output"])
|
|
# Sharded matmuls + all-reduce reorder the summation, so allow a small tolerance over the reference.
|
|
torch.testing.assert_close(ref_output, tp_output, atol=1e-3, rtol=1e-3)
|
|
|
|
def test_tensor_parallel_batch_inputs(self):
|
|
self.test_tensor_parallel_inference(batch_size=2)
|
|
|
|
def _tp_checkpoint_and_reference(self, tmp_path, world_size, sharded):
|
|
"""Write a checkpoint for the sharded loaders to read, and record its single-device output.
|
|
|
|
Returns `(checkpoint_dir, cpu_inputs, reference_output)`, or skips when the model cannot be sharded
|
|
across `world_size` ranks.
|
|
"""
|
|
if not torch.distributed.is_available():
|
|
pytest.skip("torch.distributed is not available.")
|
|
if getattr(self.model_class, "_tp_plan", None) is None:
|
|
pytest.skip("Model does not define a `_tp_plan` for tensor parallel inference.")
|
|
|
|
init_dict = self.get_init_dict()
|
|
num_heads = init_dict.get("num_attention_heads")
|
|
if num_heads is not None and num_heads % world_size != 0:
|
|
pytest.skip(f"`num_attention_heads` ({num_heads}) is not divisible by tp_degree ({world_size}).")
|
|
|
|
inputs_dict = self.get_dummy_inputs()
|
|
model = self.model_class(**init_dict).eval().to(torch_device)
|
|
with torch.no_grad():
|
|
reference = model(**inputs_dict, return_dict=False)[0].float().cpu()
|
|
|
|
checkpoint_dir = str(tmp_path / "checkpoint")
|
|
max_shard_size = int(compute_module_persistent_sizes(model)[""] * 0.75) if sharded else "5GB"
|
|
model.save_pretrained(checkpoint_dir, max_shard_size=max_shard_size)
|
|
|
|
if sharded:
|
|
index_path = os.path.join(checkpoint_dir, SAFE_WEIGHTS_INDEX_NAME)
|
|
assert os.path.exists(index_path)
|
|
expected_num_shards = calculate_expected_num_shards(index_path)
|
|
actual_num_shards = len([file for file in os.listdir(checkpoint_dir) if file.endswith(".safetensors")])
|
|
assert actual_num_shards == expected_num_shards > 1
|
|
|
|
inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
|
|
return checkpoint_dir, inputs_dict, reference
|
|
|
|
@pytest.mark.parametrize("sharded", [False, True], ids=["unsharded", "sharded"])
|
|
def test_tensor_parallel_from_pretrained(self, tmp_path, sharded):
|
|
"""`from_pretrained(..., parallel_config=...)` shards while reading and matches the single-device reference."""
|
|
world_size = 2
|
|
checkpoint_dir, inputs_dict, reference = self._tp_checkpoint_and_reference(tmp_path, world_size, sharded)
|
|
|
|
manager = mp.Manager()
|
|
return_dict = manager.dict()
|
|
mp.spawn(
|
|
_tensor_parallel_from_pretrained_worker,
|
|
args=(
|
|
world_size,
|
|
_find_free_port(),
|
|
self.model_class,
|
|
checkpoint_dir,
|
|
inputs_dict,
|
|
return_dict,
|
|
),
|
|
nprocs=world_size,
|
|
join=True,
|
|
)
|
|
assert return_dict.get("status") == "success", (
|
|
f"Tensor parallel `from_pretrained` failed: {return_dict.get('error', 'Unknown error')}"
|
|
)
|
|
|
|
# Sharded matmuls + all-reduce reorder the summation, so allow a small tolerance over the reference.
|
|
torch.testing.assert_close(reference, torch.tensor(return_dict["output"]), atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
@is_tensor_parallel
|
|
@require_torch_tpu
|
|
class TensorParallelTPUTesterMixin:
|
|
"""Same check as `TensorParallelTesterMixin`, spawning one process per TPU chip.
|
|
|
|
Run these tests in their own pytest process (e.g. `-k TensorParallelTPU`): `torch.manual_seed` or a backward pass
|
|
in an earlier test binds every TPU chip to the pytest process, and the spawned workers then fail.
|
|
"""
|
|
|
|
# Smallest slice every TPU generation supports, so the test is the same on any CI host.
|
|
tp_world_size = 4
|
|
|
|
def test_tensor_parallel_tpu_inference(self, atol=1e-3, rtol=1e-3):
|
|
from torch_tpu._internal.distributed.launchers.singlehost_wrapper import prepare_tpu_environment
|
|
from torch_tpu._internal.utils import hardware
|
|
|
|
if getattr(self.model_class, "_tp_plan", None) is None:
|
|
pytest.skip("Model does not define a `_tp_plan` for tensor parallel inference.")
|
|
|
|
world_size = self.tp_world_size
|
|
if hardware.get_tpu_device_count() < world_size:
|
|
pytest.skip(f"Needs at least {world_size} TPU chips.")
|
|
init_dict = self.get_init_dict()
|
|
num_heads = init_dict.get("num_attention_heads")
|
|
if num_heads is not None and num_heads % world_size != 0:
|
|
pytest.skip(f"`num_attention_heads` ({num_heads}) is not divisible by tp_degree ({world_size}).")
|
|
|
|
# Reference on CPU: touching the TPU here would bind every chip to this process.
|
|
inputs_dict = self.get_dummy_inputs(device="cpu")
|
|
model = self.model_class(**init_dict).eval()
|
|
with torch.no_grad():
|
|
ref_output = model(**inputs_dict, return_dict=False)[0].float()
|
|
|
|
tpu_env = ("TORCH_TPU_TOPOLOGY", "TORCH_TPU_SLICEBUILDER_ADDRESSES")
|
|
for key in tpu_env:
|
|
os.environ.pop(key, None)
|
|
prepare_tpu_environment(world_size)
|
|
return_dict = mp.Manager().dict()
|
|
try:
|
|
mp.spawn(
|
|
_tensor_parallel_worker,
|
|
args=(
|
|
world_size,
|
|
_find_free_port(),
|
|
self.model_class,
|
|
init_dict,
|
|
inputs_dict,
|
|
return_dict,
|
|
model.state_dict(),
|
|
"tpu",
|
|
),
|
|
nprocs=world_size,
|
|
join=True,
|
|
)
|
|
finally:
|
|
for key in tpu_env:
|
|
os.environ.pop(key, None)
|
|
|
|
assert return_dict.get("status") == "success", (
|
|
f"Tensor parallel inference failed: {return_dict.get('error', 'Unknown error')}"
|
|
)
|
|
tp_output = torch.tensor(return_dict["output"])
|
|
torch.testing.assert_close(ref_output, tp_output, atol=atol, rtol=rtol)
|
|
|
|
|
|
@is_tensor_parallel
|
|
@require_torch_neuron
|
|
class TensorParallelNeuronTesterMixin:
|
|
"""Same check as `TensorParallelTesterMixin`, spawning one process per Neuron core.
|
|
|
|
Run these tests in their own pytest process (e.g. `-k TensorParallelNeuron`): an earlier test that allocates a
|
|
Neuron tensor initializes the Neuron runtime in the pytest process, and the spawned workers then fail to initialize
|
|
theirs.
|
|
"""
|
|
|
|
tp_world_size = 2
|
|
|
|
def test_tensor_parallel_neuron_inference(self, atol=2e-2, rtol=2e-2):
|
|
if getattr(self.model_class, "_tp_plan", None) is None:
|
|
pytest.skip("Model does not define a `_tp_plan` for tensor parallel inference.")
|
|
|
|
world_size = self.tp_world_size
|
|
init_dict = self.get_init_dict()
|
|
num_heads = init_dict.get("num_attention_heads")
|
|
if num_heads is not None and num_heads % world_size != 0:
|
|
pytest.skip(f"`num_attention_heads` ({num_heads}) is not divisible by tp_degree ({world_size}).")
|
|
|
|
# Reference on CPU: touching the Neuron device here would break the spawned workers.
|
|
inputs_dict = self.get_dummy_inputs(device="cpu")
|
|
model = self.model_class(**init_dict).eval()
|
|
with torch.no_grad():
|
|
ref_output = model(**inputs_dict, return_dict=False)[0].float()
|
|
|
|
return_dict = mp.Manager().dict()
|
|
mp.spawn(
|
|
_tensor_parallel_worker,
|
|
args=(
|
|
world_size,
|
|
_find_free_port(),
|
|
self.model_class,
|
|
init_dict,
|
|
inputs_dict,
|
|
return_dict,
|
|
model.state_dict(),
|
|
"neuron",
|
|
),
|
|
nprocs=world_size,
|
|
join=True,
|
|
)
|
|
|
|
assert return_dict.get("status") == "success", (
|
|
f"Tensor parallel inference failed: {return_dict.get('error', 'Unknown error')}"
|
|
)
|
|
tp_output = torch.tensor(return_dict["output"])
|
|
# Neuron runs matmuls in bf16 internally, so compare with a bf16-level tolerance.
|
|
torch.testing.assert_close(ref_output, tp_output, atol=atol, rtol=rtol)
|
|
|
|
|
|
@is_context_parallel
|
|
@require_torch_multi_accelerator
|
|
class ContextParallelTesterMixin:
|
|
@pytest.mark.parametrize("cp_type", ["ulysses_degree", "ring_degree"], ids=["ulysses", "ring"])
|
|
def test_context_parallel_inference(self, cp_type, batch_size: int = 1):
|
|
if not torch.distributed.is_available():
|
|
pytest.skip("torch.distributed is not available.")
|
|
|
|
if not hasattr(self.model_class, "_cp_plan") or self.model_class._cp_plan is None:
|
|
pytest.skip("Model does not have a _cp_plan defined for context parallel inference.")
|
|
|
|
if cp_type == "ring_degree":
|
|
active_backend, _ = _AttentionBackendRegistry.get_active_backend()
|
|
if active_backend == AttentionBackendName.NATIVE:
|
|
pytest.skip("Ring attention is not supported with the native attention backend.")
|
|
|
|
world_size = 2
|
|
init_dict = self.get_init_dict()
|
|
inputs_dict = self.get_dummy_inputs(batch_size=batch_size)
|
|
|
|
# Single-GPU reference
|
|
model = self.model_class(**init_dict).eval().to(torch_device)
|
|
state_dict = {k: v.cpu() for k, v in model.state_dict().items()}
|
|
with torch.no_grad():
|
|
ref_output = model(**inputs_dict, return_dict=False)[0].cpu()
|
|
|
|
# Move all tensors to CPU for multiprocessing
|
|
inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
|
|
cp_dict = {cp_type: world_size}
|
|
|
|
# Find a free port for distributed communication
|
|
master_port = _find_free_port()
|
|
|
|
# Use multiprocessing manager for cross-process communication
|
|
manager = mp.Manager()
|
|
return_dict = manager.dict()
|
|
|
|
# Spawn worker processes
|
|
mp.spawn(
|
|
_context_parallel_worker,
|
|
args=(
|
|
world_size,
|
|
master_port,
|
|
self.model_class,
|
|
init_dict,
|
|
cp_dict,
|
|
inputs_dict,
|
|
return_dict,
|
|
None,
|
|
state_dict,
|
|
),
|
|
nprocs=world_size,
|
|
join=True,
|
|
)
|
|
|
|
assert return_dict.get("status") == "success", (
|
|
f"Context parallel inference failed: {return_dict.get('error', 'Unknown error')}"
|
|
)
|
|
|
|
cp_output = torch.tensor(return_dict["output"])
|
|
torch.testing.assert_close(ref_output, cp_output, atol=1e-4, rtol=1e-4)
|
|
|
|
@pytest.mark.parametrize("cp_type", ["ulysses_degree", "ring_degree"], ids=["ulysses", "ring"])
|
|
def test_context_parallel_batch_inputs(self, cp_type):
|
|
self.test_context_parallel_inference(cp_type, batch_size=2)
|
|
|
|
@pytest.mark.parametrize("cp_type", ["ulysses_degree", "ring_degree"], ids=["ulysses", "ring"])
|
|
def test_context_parallel_backward(self, cp_type, batch_size: int = 1):
|
|
if not torch.distributed.is_available():
|
|
pytest.skip("torch.distributed is not available.")
|
|
|
|
if not hasattr(self.model_class, "_cp_plan") or self.model_class._cp_plan is None:
|
|
pytest.skip("Model does not have a _cp_plan defined for context parallel inference.")
|
|
|
|
if cp_type == "ring_degree":
|
|
active_backend, _ = _AttentionBackendRegistry.get_active_backend()
|
|
if active_backend == AttentionBackendName.NATIVE:
|
|
pytest.skip("Ring attention is not supported with the native attention backend.")
|
|
|
|
world_size = 2
|
|
init_dict = self.get_init_dict()
|
|
inputs_dict = self.get_dummy_inputs(batch_size=batch_size)
|
|
|
|
# Move all tensors to CPU for multiprocessing
|
|
inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
|
|
cp_dict = {cp_type: world_size}
|
|
|
|
# Find a free port for distributed communication
|
|
master_port = _find_free_port()
|
|
|
|
# Use multiprocessing manager for cross-process communication
|
|
manager = mp.Manager()
|
|
return_dict = manager.dict()
|
|
|
|
# Spawn worker processes
|
|
mp.spawn(
|
|
_context_parallel_backward_worker,
|
|
args=(world_size, master_port, self.model_class, init_dict, cp_dict, inputs_dict, return_dict),
|
|
nprocs=world_size,
|
|
join=True,
|
|
)
|
|
|
|
assert return_dict.get("status") == "success", (
|
|
f"Context parallel backward pass failed: {return_dict.get('error', 'Unknown error')}"
|
|
)
|
|
assert return_dict.get("has_valid_grads"), "Context parallel backward pass did not produce valid gradients."
|
|
|
|
@pytest.mark.parametrize("cp_type", ["ulysses_degree", "ring_degree"], ids=["ulysses", "ring"])
|
|
def test_context_parallel_backward_batch_inputs(self, cp_type):
|
|
self.test_context_parallel_backward(cp_type, batch_size=2)
|
|
|
|
@pytest.mark.parametrize(
|
|
"cp_type,mesh_shape,mesh_dim_names",
|
|
[
|
|
("ring_degree", (2, 1, 1), ("ring", "ulysses", "fsdp")),
|
|
("ulysses_degree", (1, 2, 1), ("ring", "ulysses", "fsdp")),
|
|
],
|
|
ids=["ring-3d-fsdp", "ulysses-3d-fsdp"],
|
|
)
|
|
def test_context_parallel_custom_mesh(self, cp_type, mesh_shape, mesh_dim_names):
|
|
if not torch.distributed.is_available():
|
|
pytest.skip("torch.distributed is not available.")
|
|
|
|
if not hasattr(self.model_class, "_cp_plan") or self.model_class._cp_plan is None:
|
|
pytest.skip("Model does not have a _cp_plan defined for context parallel inference.")
|
|
|
|
if cp_type == "ring_degree":
|
|
active_backend, _ = _AttentionBackendRegistry.get_active_backend()
|
|
if active_backend == AttentionBackendName.NATIVE:
|
|
pytest.skip("Ring attention is not supported with the native attention backend.")
|
|
|
|
world_size = 2
|
|
init_dict = self.get_init_dict()
|
|
inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in self.get_dummy_inputs().items()}
|
|
cp_dict = {cp_type: world_size}
|
|
|
|
master_port = _find_free_port()
|
|
manager = mp.Manager()
|
|
return_dict = manager.dict()
|
|
|
|
mp.spawn(
|
|
_custom_mesh_worker,
|
|
args=(
|
|
world_size,
|
|
master_port,
|
|
self.model_class,
|
|
init_dict,
|
|
cp_dict,
|
|
mesh_shape,
|
|
mesh_dim_names,
|
|
inputs_dict,
|
|
return_dict,
|
|
),
|
|
nprocs=world_size,
|
|
join=True,
|
|
)
|
|
|
|
assert return_dict.get("status") == "success", (
|
|
f"Custom mesh context parallel inference failed: {return_dict.get('error', 'Unknown error')}"
|
|
)
|
|
|
|
|
|
@is_attention
|
|
@is_context_parallel
|
|
@require_torch_multi_accelerator
|
|
class ContextParallelAttentionBackendsTesterMixin:
|
|
unsupported_attn_backends: list[str] = []
|
|
|
|
@pytest.mark.parametrize("cp_type", ["ulysses_degree", "ring_degree"])
|
|
@pytest.mark.parametrize(
|
|
"attention_backend",
|
|
[
|
|
"native",
|
|
pytest.param(
|
|
"flash_hub",
|
|
marks=pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
|
|
),
|
|
pytest.param(
|
|
"flash_varlen_hub",
|
|
marks=pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
|
|
),
|
|
pytest.param(
|
|
"_flash_3_hub",
|
|
marks=pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
|
|
),
|
|
pytest.param(
|
|
"_flash_3_varlen_hub",
|
|
marks=pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("ulysses_anything", [True, False])
|
|
@torch.no_grad()
|
|
def test_context_parallel_attn_backend_inference(self, cp_type, attention_backend, ulysses_anything):
|
|
if not torch.distributed.is_available():
|
|
pytest.skip("torch.distributed is not available.")
|
|
|
|
if getattr(self.model_class, "_cp_plan", None) is None:
|
|
pytest.skip("Model does not have a _cp_plan defined for context parallel inference.")
|
|
|
|
if attention_backend in self.unsupported_attn_backends:
|
|
pytest.skip(f"{attention_backend} is not supported for this model.")
|
|
|
|
if cp_type == "ring_degree":
|
|
if attention_backend == AttentionBackendName.NATIVE:
|
|
pytest.skip("Skipping test because ring isn't supported with native attention backend.")
|
|
elif attention_backend in ("flash_varlen_hub", "_flash_3_varlen_hub"):
|
|
pytest.skip("`ring_degree` is not yet supported for varlen attention hub kernels.")
|
|
|
|
if ulysses_anything and "ulysses" not in cp_type:
|
|
pytest.skip("Skipping test as ulysses anything needs the ulysses degree set.")
|
|
|
|
world_size = 2
|
|
init_dict = self.get_init_dict()
|
|
inputs_dict = self.get_dummy_inputs()
|
|
|
|
# Single-GPU reference with the same attention backend (no context parallel)
|
|
model = self.model_class(**init_dict).eval().to(torch_device)
|
|
if attention_backend:
|
|
model.set_attention_backend(attention_backend)
|
|
|
|
# Copy inputs and cast model + inputs as needed
|
|
ref_inputs = inputs_dict.copy()
|
|
model, ref_inputs = _maybe_cast_to_bf16(attention_backend, model, ref_inputs)
|
|
state_dict = {k: v.cpu() for k, v in model.state_dict().items()}
|
|
with torch.no_grad():
|
|
ref_output = model(**ref_inputs, return_dict=False)[0].cpu()
|
|
|
|
# Move all tensors to CPU for multiprocessing
|
|
inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
|
|
cp_dict = {cp_type: world_size}
|
|
if ulysses_anything:
|
|
cp_dict.update({"ulysses_anything": ulysses_anything})
|
|
|
|
# Find a free port for distributed communication
|
|
master_port = _find_free_port()
|
|
|
|
# Use multiprocessing manager for cross-process communication
|
|
manager = mp.Manager()
|
|
return_dict = manager.dict()
|
|
|
|
# Spawn worker processes
|
|
mp.spawn(
|
|
_context_parallel_worker,
|
|
args=(
|
|
world_size,
|
|
master_port,
|
|
self.model_class,
|
|
init_dict,
|
|
cp_dict,
|
|
inputs_dict,
|
|
return_dict,
|
|
attention_backend,
|
|
state_dict,
|
|
),
|
|
nprocs=world_size,
|
|
join=True,
|
|
)
|
|
|
|
assert return_dict.get("status") == "success", (
|
|
f"Context parallel inference failed: {return_dict.get('error', 'Unknown error')}"
|
|
)
|
|
|
|
cp_output = torch.tensor(return_dict["output"], dtype=ref_output.dtype)
|
|
torch.testing.assert_close(ref_output, cp_output, atol=1e-2, rtol=1e-2)
|