wehub-resource-sync
406 lines
15 KiB
Python
406 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# DeepSpeed Team
|
|
"""Shared fixtures and assertions for compact AutoEP tests."""
|
|
|
|
import copy
|
|
import os
|
|
import tempfile
|
|
import traceback
|
|
from queue import Empty
|
|
|
|
import deepspeed
|
|
import deepspeed.comm as dist
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.multiprocessing as mp
|
|
|
|
from deepspeed.accelerator import get_accelerator, set_accelerator
|
|
from deepspeed.accelerator.cpu_accelerator import CPU_Accelerator
|
|
from unit.common import DEEPSPEED_TEST_TIMEOUT, get_master_port
|
|
|
|
UNSET = object()
|
|
UNSUPPORTED_LOAD_BALANCE_VALUES = [0, 0.0, 1e-3, 0.02, False, True, "1e-3", [1e-3], {"coeff": 1e-3}]
|
|
H100_TEST_ENV_VARS = ("DEEPSPEED_RUN_H100_TESTS", "DEVDS_RUN_H100_TESTS")
|
|
|
|
|
|
def h100_tests_enabled():
|
|
return any(os.environ.get(name) for name in H100_TEST_ENV_VARS)
|
|
|
|
|
|
def skip_unless_h100_tests_enabled(reason):
|
|
if not h100_tests_enabled():
|
|
pytest.skip(f"{reason}; set DEEPSPEED_RUN_H100_TESTS=1 or DEVDS_RUN_H100_TESTS=1")
|
|
|
|
|
|
class MockHFConfig:
|
|
model_type = "mixtral"
|
|
num_local_experts = 4
|
|
num_experts_per_tok = 2
|
|
hidden_size = 64
|
|
intermediate_size = 128
|
|
|
|
|
|
class MockMoEExperts(nn.Module):
|
|
|
|
def __init__(self, num_experts=4, ffn_hidden=128, hidden_size=64, intermediate_size=None):
|
|
super().__init__()
|
|
if intermediate_size is not None:
|
|
ffn_hidden = intermediate_size
|
|
self.gate_up_proj = nn.Parameter(torch.randn(num_experts, 2 * ffn_hidden, hidden_size))
|
|
self.down_proj = nn.Parameter(torch.randn(num_experts, hidden_size, ffn_hidden))
|
|
|
|
|
|
class MockMoEBlock(nn.Module):
|
|
|
|
def __init__(self, num_experts=4, ffn_hidden=128, hidden_size=64, intermediate_size=None):
|
|
super().__init__()
|
|
if intermediate_size is not None:
|
|
ffn_hidden = intermediate_size
|
|
self.gate = nn.Linear(hidden_size, num_experts, bias=False)
|
|
self.experts = MockMoEExperts(num_experts, ffn_hidden, hidden_size, intermediate_size)
|
|
self.top_k = 2
|
|
|
|
def forward(self, x):
|
|
original_shape = x.shape
|
|
hidden_states = x.reshape(-1, original_shape[-1])
|
|
scores = torch.softmax(self.gate(hidden_states), dim=-1)
|
|
top_scores, top_indices = torch.topk(scores, k=self.top_k, dim=-1)
|
|
top_scores = top_scores / top_scores.sum(dim=-1, keepdim=True)
|
|
output = torch.zeros_like(hidden_states)
|
|
|
|
for expert_idx in range(self.gate.out_features):
|
|
expert_mask = top_indices == expert_idx
|
|
if not expert_mask.any():
|
|
continue
|
|
token_indices, route_indices = expert_mask.nonzero(as_tuple=True)
|
|
expert_input = hidden_states[token_indices]
|
|
gate_up = torch.matmul(expert_input, self.experts.gate_up_proj[expert_idx].transpose(0, 1))
|
|
gate_part, up_part = gate_up.chunk(2, dim=-1)
|
|
expert_output = torch.matmul(
|
|
torch.nn.functional.silu(gate_part) * up_part, self.experts.down_proj[expert_idx].transpose(0, 1))
|
|
output[token_indices] += expert_output * top_scores[token_indices, route_indices].unsqueeze(-1)
|
|
|
|
return output.reshape(original_shape)
|
|
|
|
|
|
class MockDenseBlock(nn.Module):
|
|
|
|
def __init__(self, hidden_size=64, ffn_hidden=128):
|
|
super().__init__()
|
|
self.gate_proj = nn.Linear(hidden_size, ffn_hidden, bias=False)
|
|
self.up_proj = nn.Linear(hidden_size, ffn_hidden, bias=False)
|
|
self.down_proj = nn.Linear(ffn_hidden, hidden_size, bias=False)
|
|
|
|
|
|
class MockMoETransformer(nn.Module):
|
|
|
|
def __init__(self, num_layers=2, num_experts=4, hidden_size=64, intermediate_size=128, moe_every_n=1):
|
|
super().__init__()
|
|
self.config = MockHFConfig()
|
|
self.config.num_local_experts = num_experts
|
|
self.config.hidden_size = hidden_size
|
|
self.config.intermediate_size = intermediate_size
|
|
self.model = nn.Module()
|
|
self.model.layers = nn.ModuleList([
|
|
self._make_layer(layer_idx, num_experts, hidden_size, intermediate_size, moe_every_n)
|
|
for layer_idx in range(num_layers)
|
|
])
|
|
self.lm_head = nn.Linear(hidden_size, 100)
|
|
|
|
@staticmethod
|
|
def _make_layer(layer_idx, num_experts, hidden_size, intermediate_size, moe_every_n):
|
|
layer = nn.Module()
|
|
layer.self_attn = nn.MultiheadAttention(hidden_size, 1, batch_first=True)
|
|
if layer_idx % moe_every_n == 0:
|
|
layer.mlp = MockMoEBlock(num_experts, intermediate_size, hidden_size)
|
|
else:
|
|
layer.mlp = MockDenseBlock(hidden_size, intermediate_size)
|
|
layer.input_layernorm = nn.LayerNorm(hidden_size)
|
|
layer.post_attention_layernorm = nn.LayerNorm(hidden_size)
|
|
return layer
|
|
|
|
def forward(self, x):
|
|
for layer_module in self.model.layers:
|
|
residual = x
|
|
x = layer_module.input_layernorm(x)
|
|
x, _ = layer_module.self_attn(x, x, x)
|
|
x = residual + x
|
|
residual = x
|
|
x = layer_module.post_attention_layernorm(x)
|
|
x = residual + layer_module.mlp(x)
|
|
return self.lm_head(x)
|
|
|
|
|
|
class MockMoEOnlyTransformer(nn.Module):
|
|
|
|
def __init__(self, num_layers=2, num_experts=4, hidden_size=64, intermediate_size=128, moe_every_n=1):
|
|
super().__init__()
|
|
self.config = MockHFConfig()
|
|
self.config.num_local_experts = num_experts
|
|
self.config.hidden_size = hidden_size
|
|
self.config.intermediate_size = intermediate_size
|
|
self.model = nn.Module()
|
|
self.model.layers = nn.ModuleList([
|
|
self._make_layer(layer_idx, num_experts, hidden_size, intermediate_size, moe_every_n)
|
|
for layer_idx in range(num_layers)
|
|
])
|
|
self.lm_head = nn.Linear(hidden_size, 100, bias=False)
|
|
|
|
@staticmethod
|
|
def _make_layer(layer_idx, num_experts, hidden_size, intermediate_size, moe_every_n):
|
|
layer = nn.Module()
|
|
layer.dense = nn.Linear(hidden_size, hidden_size, bias=False)
|
|
if layer_idx % moe_every_n == 0:
|
|
layer.mlp = MockMoEBlock(num_experts, intermediate_size, hidden_size)
|
|
else:
|
|
layer.mlp = MockDenseBlock(hidden_size, intermediate_size)
|
|
layer.input_layernorm = nn.LayerNorm(hidden_size)
|
|
layer.post_attention_layernorm = nn.LayerNorm(hidden_size)
|
|
return layer
|
|
|
|
def forward(self, x):
|
|
for layer_module in self.model.layers:
|
|
residual = x
|
|
x = layer_module.input_layernorm(x)
|
|
x = residual + layer_module.dense(x)
|
|
residual = x
|
|
x = layer_module.post_attention_layernorm(x)
|
|
x = residual + layer_module.mlp(x)
|
|
return self.lm_head(x)
|
|
|
|
|
|
def assert_load_balance_coeff_rejection_message(exc: BaseException, value: object) -> None:
|
|
text = str(exc)
|
|
for needle in ("load_balance_coeff", "expert_bias", "not supported", "null", "omit"):
|
|
assert needle in text
|
|
assert repr(value) in text
|
|
|
|
|
|
def mixed_precision_config():
|
|
accelerator = get_accelerator()
|
|
if accelerator.is_fp16_supported() and accelerator.device_name() != "cpu":
|
|
return {"fp16": {"enabled": True, "initial_scale_power": 8}}
|
|
if accelerator.is_bf16_supported():
|
|
return {"bf16": {"enabled": True}}
|
|
if accelerator.is_fp16_supported():
|
|
return {"fp16": {"enabled": True, "initial_scale_power": 8}}
|
|
pytest.skip("AutoEP tests require fp16 or bf16 support")
|
|
|
|
|
|
def make_autoep_config(zero_stage=0, ep_size=1, load_balance_coeff=UNSET, mixed_precision=True):
|
|
config = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1e-4
|
|
},
|
|
},
|
|
"expert_parallel": {
|
|
"enabled": True,
|
|
"autoep_size": ep_size,
|
|
"preset_model": "mixtral",
|
|
"use_grouped_mm": False,
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
},
|
|
}
|
|
if get_accelerator().device_name() == "cpu":
|
|
config["optimizer"]["params"]["torch_adam"] = True
|
|
if mixed_precision:
|
|
config.update(mixed_precision_config())
|
|
if load_balance_coeff is not UNSET:
|
|
config["expert_parallel"]["load_balance_coeff"] = load_balance_coeff
|
|
return config
|
|
|
|
|
|
def make_autoep_integration_config(zero_stage=0, ep_size=2):
|
|
return make_autoep_config(zero_stage=zero_stage, ep_size=ep_size, mixed_precision=False)
|
|
|
|
|
|
def seed_everything(seed=1234):
|
|
torch.manual_seed(seed)
|
|
get_accelerator().manual_seed(seed)
|
|
get_accelerator().manual_seed_all(seed)
|
|
|
|
|
|
def engine_input_dtype(engine):
|
|
if engine.bfloat16_enabled():
|
|
return torch.bfloat16
|
|
if engine.fp16_enabled():
|
|
return torch.float16
|
|
return torch.float32
|
|
|
|
|
|
def init_autoep_engine(ep_size=1, zero_stage=0, load_balance_coeff=UNSET):
|
|
seed_everything(42)
|
|
engine, _, _, _ = deepspeed.initialize(
|
|
model=MockMoETransformer(),
|
|
config=make_autoep_config(zero_stage=zero_stage, ep_size=ep_size, load_balance_coeff=load_balance_coeff),
|
|
)
|
|
return engine
|
|
|
|
|
|
def run_training_steps(engine, num_steps=3, seq_len=8, hidden_dim=64):
|
|
losses = []
|
|
grad_norms = []
|
|
for _ in range(num_steps):
|
|
x = torch.randn(1, seq_len, hidden_dim, device=engine.device)
|
|
loss = engine(x).mean()
|
|
engine.backward(loss)
|
|
|
|
total_norm = 0.0
|
|
for param in engine.module.parameters():
|
|
if param.grad is not None:
|
|
total_norm += param.grad.data.float().norm(2).item()**2
|
|
grad_norms.append(total_norm**0.5)
|
|
engine.step()
|
|
losses.append(loss.item())
|
|
return losses, grad_norms
|
|
|
|
|
|
def tiny_causal_lm_inputs():
|
|
input_ids = torch.tensor([[1, 5, 7, 9, 11]], dtype=torch.long)
|
|
return input_ids, input_ids.clone()
|
|
|
|
|
|
def state_matched_models(model_cls, config):
|
|
native_model = model_cls(config)
|
|
autoep_model = model_cls(config)
|
|
autoep_model.load_state_dict(copy.deepcopy(native_model.state_dict()))
|
|
return native_model, autoep_model
|
|
|
|
|
|
def replace_autoep_layers(model, preset_model, expected_count=1, **config_overrides):
|
|
from deepspeed.module_inject.auto_ep import AutoEP
|
|
from deepspeed.module_inject.auto_ep_config import parse_autoep_config
|
|
|
|
config = {
|
|
"enabled": True,
|
|
"autoep_size": 1,
|
|
"preset_model": preset_model,
|
|
"use_grouped_mm": False,
|
|
**config_overrides,
|
|
}
|
|
auto_ep = AutoEP(model, parse_autoep_config(config))
|
|
specs = auto_ep.ep_parser()
|
|
assert len(specs) == expected_count
|
|
for spec in specs:
|
|
auto_ep.replace_moe_layer(spec, ep_size=1, ep_rank=0)
|
|
return specs
|
|
|
|
|
|
def assert_causal_lm_outputs_close(native_model,
|
|
autoep_model,
|
|
*,
|
|
output_router_logits=False,
|
|
compare_router_logits=False,
|
|
compare_aux_loss=False,
|
|
compare_logits=True,
|
|
rtol=1e-5,
|
|
atol=1e-6):
|
|
input_ids, labels = tiny_causal_lm_inputs()
|
|
native_model.eval()
|
|
autoep_model.eval()
|
|
with torch.no_grad():
|
|
native_outputs = native_model(input_ids=input_ids, labels=labels, output_router_logits=output_router_logits)
|
|
autoep_outputs = autoep_model(input_ids=input_ids, labels=labels, output_router_logits=output_router_logits)
|
|
|
|
if compare_router_logits:
|
|
assert autoep_outputs.router_logits
|
|
torch.testing.assert_close(autoep_outputs.router_logits[0],
|
|
native_outputs.router_logits[0],
|
|
rtol=rtol,
|
|
atol=atol)
|
|
if compare_aux_loss:
|
|
assert autoep_outputs.aux_loss is not None
|
|
torch.testing.assert_close(autoep_outputs.aux_loss, native_outputs.aux_loss, rtol=rtol, atol=atol)
|
|
if compare_logits:
|
|
torch.testing.assert_close(autoep_outputs.logits, native_outputs.logits, rtol=rtol, atol=atol)
|
|
torch.testing.assert_close(autoep_outputs.loss, native_outputs.loss, rtol=rtol, atol=atol)
|
|
|
|
|
|
def skip_unless_transformers_has(transformers, *names, min_version=None, reason="AutoEP coverage"):
|
|
from packaging.version import Version
|
|
|
|
if min_version is not None and Version(transformers.__version__) < Version(min_version):
|
|
pytest.skip(f"{reason} requires Transformers >= {min_version}")
|
|
missing = [name for name in names if not hasattr(transformers, name)]
|
|
if missing:
|
|
pytest.skip(f"Installed transformers does not expose required classes: {missing}")
|
|
|
|
|
|
def tiny_mixtral_config(transformers):
|
|
return transformers.MixtralConfig(
|
|
vocab_size=64,
|
|
hidden_size=32,
|
|
intermediate_size=64,
|
|
num_hidden_layers=1,
|
|
num_attention_heads=4,
|
|
num_key_value_heads=2,
|
|
max_position_embeddings=64,
|
|
num_local_experts=4,
|
|
num_experts_per_tok=2,
|
|
output_router_logits=True,
|
|
tie_word_embeddings=False,
|
|
use_cache=False,
|
|
)
|
|
|
|
|
|
def _cpu_gloo_worker_entry(rank, world_size, init_method, master_port, worker, shared_tmpdir, error_queue):
|
|
set_accelerator(CPU_Accelerator())
|
|
os.environ["MASTER_ADDR"] = "127.0.0.1"
|
|
os.environ["MASTER_PORT"] = master_port
|
|
os.environ["LOCAL_RANK"] = str(rank)
|
|
os.environ["RANK"] = str(rank)
|
|
os.environ["LOCAL_SIZE"] = str(world_size)
|
|
os.environ["WORLD_SIZE"] = str(world_size)
|
|
os.environ.pop("NCCL_DEBUG", None)
|
|
|
|
try:
|
|
deepspeed.init_distributed(dist_backend="gloo", init_method=init_method, rank=rank, world_size=world_size)
|
|
worker(rank, world_size, shared_tmpdir)
|
|
except BaseException:
|
|
error_queue.put(traceback.format_exc())
|
|
raise
|
|
finally:
|
|
if dist.is_initialized():
|
|
dist.destroy_process_group()
|
|
|
|
|
|
def run_cpu_gloo_test(worker, tmpdir, *, world_size=4, timeout=DEEPSPEED_TEST_TIMEOUT):
|
|
"""Run a small CPU/Gloo distributed test without requiring visible GPU devices."""
|
|
ctx = mp.get_context("spawn")
|
|
error_queue = ctx.Queue()
|
|
with tempfile.NamedTemporaryFile(delete=False, dir=str(tmpdir), suffix="_filestore") as fp:
|
|
init_method = f"file://{fp.name}"
|
|
master_port = get_master_port()
|
|
shared_tmpdir = str(tmpdir)
|
|
processes = [
|
|
ctx.Process(target=_cpu_gloo_worker_entry,
|
|
args=(rank, world_size, init_method, master_port, worker, shared_tmpdir, error_queue))
|
|
for rank in range(world_size)
|
|
]
|
|
for process in processes:
|
|
process.start()
|
|
for process in processes:
|
|
process.join(timeout)
|
|
for process in processes:
|
|
if process.is_alive():
|
|
process.terminate()
|
|
pytest.fail(f"CPU/Gloo worker {process.pid} timed out after {timeout}s", pytrace=False)
|
|
errors = []
|
|
while True:
|
|
try:
|
|
errors.append(error_queue.get_nowait())
|
|
except Empty:
|
|
break
|
|
failed = [process for process in processes if process.exitcode]
|
|
if errors:
|
|
pytest.fail("\n".join(errors), pytrace=False)
|
|
if failed:
|
|
pytest.fail("CPU/Gloo worker failures: " + ", ".join(str(process.exitcode) for process in failed),
|
|
pytrace=False)
|