# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team from copy import deepcopy import os import random import numpy as np import torch import deepspeed import deepspeed.comm as dist from deepspeed.accelerator import get_accelerator from deepspeed.runtime.zero import GatheredParameters from unit.simple_model import SimpleModel from unit.common import allclose_on_all_ranks def compare_loss(self, config, dtype, iteration=5, hidden_dim_override=None): hidden_dim = hidden_dim_override if hidden_dim_override is not None else 10 # the default tolerances of torch.testing.assert_close are too small RTOL = 5e-1 ATOL = 1e-2 # Use a fixed seed for determinism. We don't use the @enable_determinism decorator # because it also sets torch.use_deterministic_algorithms(True), which seems # incompatible with torch.compile() in test environments. # Might be related to https://github.com/pytorch/pytorch/issues/159855 local_rank = int(os.getenv("LOCAL_RANK", "0")) seed = 123 + local_rank random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) get_accelerator().manual_seed(seed) get_accelerator().manual_seed_all(seed) device = torch.device(get_accelerator().current_device_name()) model = SimpleModel(hidden_dim) i = get_accelerator().current_device() baseline_model = deepcopy(model) baseline_config = deepcopy(config) baseline_config["zero_optimization"]["stage"] = 0 baseline_config["zero_optimization"]["offload_optimizer"] = {} baseline_engine, baseline_optimizer, _, _ = deepspeed.initialize(config=baseline_config, model=baseline_model, model_parameters=baseline_model.parameters()) if config["zero_optimization"]["stage"] == 3: with deepspeed.zero.Init(config_dict_or_path=config): target_model = SimpleModel(hidden_dim) with GatheredParameters(target_model.parameters(), modifier_rank=0): for p1, p2 in zip(target_model.parameters(), model.parameters()): p1.data.copy_(p2.data) else: target_model = deepcopy(model) target_engine, target_optimizer, _, _ = deepspeed.initialize(config=config, model=target_model, model_parameters=target_model.parameters()) target_engine.compile() train_batch_size = config["train_micro_batch_size_per_gpu"] xs = [torch.randn(train_batch_size, hidden_dim, device=device, dtype=dtype) for _ in range(iteration)] ys = [torch.randn_like(x) for x in xs] for x, y in zip(xs, ys): baseline_loss = baseline_engine(x, y) target_loss = target_engine(x, y) allclose_on_all_ranks(baseline_loss, target_loss, "Loss values are not close.", rtol=RTOL, atol=ATOL) baseline_engine.backward(baseline_loss) target_engine.backward(target_loss) baseline_engine.step() target_engine.step() with GatheredParameters(target_engine.parameters()): for p1, p2 in zip(baseline_engine.parameters(), target_engine.parameters()): allclose_on_all_ranks(p1, p2, "Parameters are not equal.", rtol=RTOL, atol=ATOL) baseline_engine.destroy() target_engine.destroy() def compare_sp_loss(self, config, sp_size, iterations=3): """ Compare AutoSP compiled model loss against a compiled Ulysses SP model (ground truth). Both engines are trained in lockstep. After all training steps the final-step losses are compared. """ import torch.nn.functional as F from transformers import AutoModelForCausalLM, AutoConfig from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS from deepspeed.compile import constants as autosp_constants from deepspeed.compile.custom_ops.sp_dp_registry import populate_registry, get_group from deepspeed.sequence.layer import DistributedAttention RTOL, ATOL = 0.1, 0.01 model_name = 'hf-internal-testing/tiny-random-LlamaForCausalLM' seq_length = 64 torch.manual_seed(42) get_accelerator().manual_seed_all(42) device = torch.device(get_accelerator().current_device_name()) model_config = AutoConfig.from_pretrained(model_name) model_config._attn_implementation = "sdpa" base_model = AutoModelForCausalLM.from_pretrained(model_name, config=model_config) vocab_size = model_config.vocab_size # Set up SP/DP process groups (shared by both Ulysses and AutoSP). dp_size = dist.get_world_size() // sp_size populate_registry(sp_size, dp_size) # The DP-rank index selects which SP group the current rank belongs to. sp_group = get_group(dist.get_rank() // sp_size) sp_rank = dist.get_rank() % sp_size chunk = seq_length // sp_size # Build a DistributedAttention wrapper that mirrors distributed_attention.py. # Registered under a unique key so the model's "sdpa" slot stays untouched — # AutoSP's graph pass can therefore find F.scaled_dot_product_attention nodes. def _sdpa_inner(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, scale=None): # DistributedAttention delivers tensors in [b, s, n, h]; SDPA wants [b, n, s, h]. out = F.scaled_dot_product_attention(q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3), dropout_p=dropout_p, is_causal=is_causal, scale=scale) return out.permute(0, 2, 1, 3) _dist_attn = DistributedAttention(_sdpa_inner, sp_group, scatter_idx=2, gather_idx=1) def _ulysses_attn_forward(module, query_states, key_states, value_states, attention_mask, scaling=None, dropout=0.0, is_causal=False, **kwargs): q = query_states.transpose(1, 2).contiguous() k = key_states.transpose(1, 2).contiguous() v = value_states.transpose(1, 2).contiguous() out = _dist_attn(q, k, v, batch_dim_idx=0, dropout_p=dropout, is_causal=is_causal, scale=scaling) return out, None ALL_ATTENTION_FUNCTIONS["ulyssess"] = _ulysses_attn_forward # Ulysses baseline: regular torch.compile, no deepcompile or autosp pass. ulysses_config = deepcopy(config) ulysses_config.pop("compile", None) ulysses_model = deepcopy(base_model) ulysses_model.config._attn_implementation = "ulyssess" ulysses_engine, _, _, _ = deepspeed.initialize(config=ulysses_config, model=ulysses_model, model_parameters=ulysses_model.parameters()) ulysses_engine.compile() # AutoSP model: sdpa so the autosp pass can find F.scaled_dot_product_attention. # dynamic=True ensures all shape dimensions are treated symbolically so the autosp # pass can correctly shard the sequence dimension for all dtypes including fp16/bf16. autosp_model = deepcopy(base_model) autosp_engine, _, _, _ = deepspeed.initialize(config=config, model=autosp_model, model_parameters=autosp_model.parameters()) autosp_engine.compile(compile_kwargs={"dynamic": True}) # Train both engines in lockstep; compare the losses at the final step. ul_loss = autosp_loss = None for i in range(iterations): torch.manual_seed(42 + i) full_ids = torch.randint(0, vocab_size, (1, seq_length), device=device) # Ulysses: each rank processes its own shard. shard_ids = full_ids[:, sp_rank * chunk:(sp_rank + 1) * chunk] shard_pos = torch.arange(sp_rank * chunk, (sp_rank + 1) * chunk, device=device).unsqueeze(0) shard_mask = torch.ones(1, chunk, device=device, dtype=torch.long) ul_out = ulysses_engine(input_ids=shard_ids, labels=shard_ids, position_ids=shard_pos, attention_mask=shard_mask) # Average per-shard losses across SP ranks to get the full-sequence loss. ul_loss = ul_out.loss.clone() dist.all_reduce(ul_loss, group=sp_group) ul_loss = ul_loss / sp_size # AutoSP: full sequence. dynamic=True makes all shapes symbolic, so mark_dynamic # is not needed; only the tag attributes that the autosp pass uses are set here. autosp_ids = full_ids.clone() autosp_lbl = autosp_ids.clone() autosp_pos = torch.arange(seq_length, device=device).unsqueeze(0) autosp_msk = torch.ones(1, seq_length, device=device, dtype=torch.long) autosp_ids.tag = autosp_constants.AUTOSP_INPUT_ID_KEY autosp_lbl.tag = autosp_constants.AUTOSP_LABEL_ID_KEY autosp_pos.tag = autosp_constants.AUTOSP_POSITION_ID_KEY autosp_out = autosp_engine(input_ids=autosp_ids, labels=autosp_lbl, position_ids=autosp_pos, attention_mask=autosp_msk) autosp_loss = autosp_out.loss ulysses_engine.backward(ul_out.loss) ulysses_engine.step() autosp_engine.backward(autosp_loss) autosp_engine.step() allclose_on_all_ranks(autosp_loss, ul_loss, "AutoSP and Ulysses losses are not close.", rtol=RTOL, atol=ATOL) ulysses_engine.destroy() del ALL_ATTENTION_FUNCTIONS["ulyssess"] autosp_engine.destroy() def create_gm_nodes(batch_size: int = 1, seq_len: int = 16): """ Load a tiny LlamaForCausalLM, tag inputs with AutoSP keys, mark the sequence dimension dynamic, and capture the torch-fx GraphModule via a custom torch.compile backend. The returned gm is identical to what the autosp pass receives during training: placeholder nodes carry tensor_dict tags and meta['val'] shapes are symbolic (SymInt) in the sequence dimension. Returns: gm – GraphModule with fully populated node metadata inputs – (input_ids, labels, position_ids) used for tracing """ from transformers import AutoModelForCausalLM, AutoConfig from deepspeed.compile import constants # Each call needs a clean dynamo state; without this, the recompile_limit # (default 8) is exhausted across tests and the backend is never invoked. torch._dynamo.reset() model_name = 'hf-internal-testing/tiny-random-LlamaForCausalLM' model_config = AutoConfig.from_pretrained(model_name) model_config._attn_implementation = "sdpa" model = AutoModelForCausalLM.from_pretrained(model_name, config=model_config) model.eval() vocab_size = model_config.vocab_size input_ids = torch.randint(0, vocab_size, (batch_size, seq_len)) labels = torch.randint(0, vocab_size, (batch_size, seq_len)) position_ids = torch.arange(seq_len).unsqueeze(0) # dynamo propagates Python tensor attributes into node.meta['tensor_dict']; # find_node_by_tag relies on this to identify the AutoSP input nodes. input_ids.tag = constants.AUTOSP_INPUT_ID_KEY labels.tag = constants.AUTOSP_LABEL_ID_KEY position_ids.tag = constants.AUTOSP_POSITION_ID_KEY # Marking the sequence dim dynamic causes dynamo to emit a SymInt placeholder # node and store symbolic shapes in node.meta['val'], which shard_tensor_node # needs to locate the sequence-length symbol in the graph. torch._dynamo.decorators.mark_dynamic(input_ids, 1) torch._dynamo.decorators.mark_dynamic(labels, 1) torch._dynamo.decorators.mark_dynamic(position_ids, 1) captured_gm = [None] def _capture_backend(gm, example_inputs): if captured_gm[0] is None: captured_gm[0] = gm return gm compiled = torch.compile(model, backend=_capture_backend, dynamic=True) with torch.no_grad(): compiled(input_ids=input_ids, labels=labels, position_ids=position_ids) assert captured_gm[0] is not None, "Capture backend was never invoked — graph capture failed" return captured_gm[0], (input_ids, labels, position_ids) def find_sym_seq_node(gm): """ Return the SymInt placeholder node for the sequence-length dimension of input_ids, or None if it cannot be found. """ from deepspeed.compile.util import get_input_id_node from deepspeed.compile.fx import get_node_shape_meta, find_node_by_name input_ids_node = get_input_id_node(gm) val = get_node_shape_meta(input_ids_node) seq_symint = val.shape[1] return find_node_by_name(gm, str(seq_symint))