wehub/deepspeedai--deepspeed
1
0
Fork 0
Code Issues Pull Requests Actions 12 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/README.md
+9
View File
@@ -0,0 +1,9 @@
# DeepSpeed Examples
If you are looking for examples using DeepSpeed please see the following resources:
1. [DeepSpeedExamples](https://github.com/deepspeedai/DeepSpeedExamples)
2. [Megatron-DeepSpeed](https://github.com/deepspeedai/Megatron-DeepSpeed)
3. [DeepSpeed + AzureML](https://github.com/Azure/azureml-examples/tree/main/v1/python-sdk/workflows/train/deepspeed)
4. [DeepSpeed + Hugging Face Transformers Integration](https://huggingface.co/docs/transformers/deepspeed)
5. [DeepSpeed + PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.utilities.deepspeed.html)
examples/sdma_allgather/README.md
+136
View File
@@ -0,0 +1,136 @@
# SDMA-Accelerated ZeRO-3 on AMD GPUs
## Motivation
ZeRO-3 reconstructs each layer with an AllGather right before its forward / backward
pass, and DeepSpeed's `PartitionedParameterCoordinator` prefetches these AllGathers
on a separate stream so that collective and compute can overlap *in time*. In
practice the time-overlap is already quite good for typical ZeRO-3 workloads.
What's left is **resource** overlap. On AMD GPUs, RCCL AllGather kernels execute
on the same compute units (CUs) that GEMM and attention run on, and tend to share
wavefront slots, LDS, register file and HBM bandwidth with concurrent compute.
Even when the time-overlap schedule looks near-perfect, the effective compute
throughput during the overlap window can stay below peak because the two workloads
sit on the same physical hardware.
AMD MI300X / MI325X / MI355X contain dedicated **System DMA (SDMA)** copy engines —
independent hardware queues that move data between GPUs over XGMI without using
the CU array. Routing ZeRO-3's AllGather through SDMA instead of CU-based RCCL
kernels lets collective traffic and compute run on physically separate engines,
leaving CUs largely free for GEMM / attention during the overlap window. In
workloads where overlap is a meaningful bottleneck this can translate into
end-to-end step-time gains (workload-dependent; see the verified results table
below).
See [RFC #7884](https://github.com/deepspeedai/DeepSpeed/issues/7884) for the
longer design rationale and discussion.
## Overview
End-to-end example for the SDMA fast-path inside
`TorchBackend.all_gather_into_tensor`. When the runtime is AMD/ROCm,
the [`mori`](https://github.com/ROCm/mori) package is importable, and the user opts in via
`DS_SDMA_ALLGATHER=1`, `deepspeed.comm` acquires the SDMA backend at
`init_distributed()` time and routes WORLD-group
`all_gather_into_tensor` calls through `mori_cpp.AllGatherIntoTensor`
(intra-node SDMA copy on MI300). RCCL/NCCL is used as the fallback on
any condition that makes the SDMA path unsafe (user did not opt in,
non-WORLD process group, shard larger than the transit buffer,
unsupported dtype, init failure).
This means:
- No `ds_config` knob — control is a single env var. Works out of the
box for ZeRO-3 (sequential and coalesced prefetch paths both benefit).
- No source modifications in `partition_parameters.py`: ZeRO-3 just calls
`dist.allgather_fn`, which lands on the backend's
`all_gather_into_tensor`.
- Sub-group allgathers (e.g. when ZeRO is initialised with a non-WORLD
data-parallel group, or with a secondary zero-param group) are routed
through RCCL/NCCL automatically, since the SDMA backend is bound to
WORLD.
- Even when mori is installed, the SDMA path stays off unless the user
sets `DS_SDMA_ALLGATHER=1`, so users keep explicit control over a
hardware-specific fast-path.
## Environment variables
| Var | Purpose |
|---|---|
| `DS_SDMA_ALLGATHER=1` | **Opt-in switch.** Required to enable the SDMA fast-path; default is off even when mori is installed. When set, `MORI_ENABLE_SDMA=1` is auto-exported on your behalf so mori allocates the uncached transit buffers the SDMA kernel needs. |
| `DS_SDMA_ALLGATHER_MAX_NUMEL=N` | Transit buffer size in elements (default 64M = 256 MiB per-rank input, ~2 GiB output on 8 ranks). Calls larger than this fall back to RCCL/NCCL. |
| `MORI_ENABLE_SDMA=1` | mori's own knob for uncached transit buffers; normally set automatically by DeepSpeed when `DS_SDMA_ALLGATHER=1`. Export it explicitly only if you want to override or pre-set it. |
The `run_*_sdma_on.sh` scripts export `DS_SDMA_ALLGATHER=1`; the
`run_*_sdma_off.sh` scripts leave it unset (default). Both variants
share the same `ds_config_zero3.json` — the SDMA decision is made
entirely by env vars.
## Verified results on 8x MI300X
| | GPT-7B-ish | Qwen3-32B |
|---|---|---|
| trainer | `train_zero3.py` | `train_qwen3_zero3.py` |
| seq / micro batch | 2048 / 1 | 1024 / 1 |
| dataset | wikitext-2-raw-v1 | wikitext-103-raw-v1 (10 %) |
| measured / warmup steps | 100 / 10 | 100 / 10 |
| **SDMA off (RCCL)** | 697.7 ms / step (11.6 samples/s) | 1402.5 ms / step (5841 tok/s) |
| **SDMA on (this PR)** | **622.0 ms / step (13.0 samples/s)** | **1263.2 ms / step (6486 tok/s)** |
| **gain** | **+10.85 %** | **+9.93 %** |
| peak mem (rank 0) | 12.12 GB, unchanged off ↔ on | 96.45 GB, unchanged off ↔ on |
The Qwen3-32B number is averaged over two fresh rounds; per-round delta
was +10.85 % and +9.92 %, with 0.29 % run-to-run variance on the off
baseline, so the gap is well outside per-step jitter (~1.5–2.7 %).
Speedup is workload-dependent — gains shrink (or invert) when allgather
cannot be overlapped with compute (e.g. very small payloads, or
`overlap_comm=false`).
### Loss curves match across off ↔ on (2000-step runs)
A long-horizon sanity check on each demo confirms the SDMA path
introduces no numerical drift: 2000 training steps on the same wikitext
shuffle, off vs on traces overlap throughout. Both trainers use the
standard "concat the corpus + slice into fixed `seq_length` chunks"
pattern, so every sample has the same number of real tokens and per-step
loss has no variance from padding fraction. Bucketed mean |off − on|
over the full 2000 steps is ≤ **0.026** on GPT and ≤ **0.048** on Qwen3,
well below natural per-step jitter.
![GPT-7B-ish — training loss vs step, SDMA off vs on, 2000 steps](images/loss_gpt_2k.png)
![Qwen3-32B — training loss vs step, SDMA off vs on, 2000 steps](images/loss_qwen3_2k.png)
## Reproduction
```bash
cd examples/sdma_allgather
# Demo 1 — GPT-7B-ish, ~minute run, no HF download
bash run_gpt_sdma_off.sh # default (DS_SDMA_ALLGATHER unset), RCCL baseline
bash run_gpt_sdma_on.sh # DS_SDMA_ALLGATHER=1, SDMA fast-path -> +10.85 %
# Demo 2 — Qwen3-32B, ~few-minute run, weight-free (random init via from_config)
bash run_qwen3_sdma_off.sh # ~1402 ms / step
bash run_qwen3_sdma_on.sh # ~1263 ms / step -> +9.93 %
```
Override knobs via env vars: `SEQ_LEN`, `BATCH_SIZE`, `NUM_STEPS`,
`WARMUP_STEPS`, `NUM_GPUS`, `MODEL`, `DS_CONFIG`.
## Files
```
ds_config_zero3.json single shared ZeRO-3 + bf16 + DS-default buckets config
run_gpt_sdma_off.sh GPT-7B-ish + ZeRO-3, SDMA disabled via env var
run_gpt_sdma_on.sh GPT-7B-ish + ZeRO-3, SDMA enabled via env var
run_qwen3_sdma_off.sh Qwen3-32B + ZeRO-3, SDMA disabled via env var
run_qwen3_sdma_on.sh Qwen3-32B + ZeRO-3, SDMA enabled via env var
test_sdma_allgather_zero3.py unit test exercising the transparent SDMA path
train_qwen3_zero3.py Qwen3 trainer (self-contained, wikitext)
train_zero3.py GPT trainer
images/loss_gpt_2k.png GPT loss curve, off vs on, 2000 steps
images/loss_qwen3_2k.png Qwen3-32B loss curve, off vs on, 2000 steps
```
examples/sdma_allgather/ds_config_zero3.json
+42
View File
@@ -0,0 +1,42 @@
{
"train_micro_batch_size_per_gpu": 1,
"gradient_accumulation_steps": 1,
"steps_per_print": 10,
"optimizer": {
"type": "AdamW",
"params": {
"lr": 1e-5,
"betas": [0.9, 0.999],
"eps": 1e-8,
"weight_decay": 0.01
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": 0,
"warmup_max_lr": 1e-5,
"warmup_num_steps": 10
}
},
"gradient_clipping": 1.0,
"bf16": {
"enabled": true
},
"zero_optimization": {
"stage": 3,
"allgather_partitions": true,
"allgather_bucket_size": 5e7,
"overlap_comm": true,
"reduce_scatter": true,
"reduce_bucket_size": 5e7,
"contiguous_gradients": true,
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_prefetch_bucket_size": 5e7,
"stage3_param_persistence_threshold": 1e5,
"stage3_gather_16bit_weights_on_model_save": true,
"sub_group_size": 1e12
},
"wall_clock_breakdown": false
}
examples/sdma_allgather/images/loss_gpt_2k.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 123 KiB

examples/sdma_allgather/images/loss_qwen3_2k.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 143 KiB

examples/sdma_allgather/run_gpt_sdma_off.sh
Executable
+12
View File
@@ -0,0 +1,12 @@
#!/bin/bash
# GPT-7B-ish + ZeRO-3 baseline (RCCL allgather).
# Default: the SDMA fast-path inside deepspeed.comm stays off unless the
# user explicitly sets DS_SDMA_ALLGATHER=1, so this script simply doesn't
# export it.
SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)
deepspeed --num_gpus 8 "${SCRIPT_DIR}/train_zero3.py" \
--deepspeed \
--deepspeed_config "${SCRIPT_DIR}/ds_config_zero3.json" \
--data_mode wikitext2 \
--train_steps 100
examples/sdma_allgather/run_gpt_sdma_on.sh
Executable
+17
View File
@@ -0,0 +1,17 @@
#!/bin/bash
# GPT-7B-ish + ZeRO-3 with the SDMA fast-path opted in.
#
# DS_SDMA_ALLGATHER=1 is the single opt-in switch. When set,
# deepspeed.comm's TorchBackend tries to bring up the mori SDMA backend
# at init time and routes WORLD-group all_gather_into_tensor through it.
# Mori's MORI_ENABLE_SDMA=1 is auto-exported on the user's behalf when
# DS_SDMA_ALLGATHER=1 is set, so users normally don't need to touch it.
# Without DS_SDMA_ALLGATHER=1, even an mori-installed run stays on RCCL.
export DS_SDMA_ALLGATHER=1
SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)
deepspeed --num_gpus 8 "${SCRIPT_DIR}/train_zero3.py" \
--deepspeed \
--deepspeed_config "${SCRIPT_DIR}/ds_config_zero3.json" \
--data_mode wikitext2 \
--train_steps 100
examples/sdma_allgather/run_qwen3_sdma_off.sh
Executable
+34
View File
@@ -0,0 +1,34 @@
#!/bin/bash
# Qwen3-32B + DeepSpeed ZeRO-3 baseline (RCCL allgather).
#
# Default: deepspeed.comm's SDMA fast-path stays off unless the user
# explicitly sets DS_SDMA_ALLGATHER=1, so this script simply doesn't
# export it and pairs cleanly with run_qwen3_sdma_on.sh (same ds_config;
# only env vars differ) for the A/B.
#
# model : Qwen/Qwen3-32B (full 64 layers, BF16, eager attention)
# data : wikitext-103-raw-v1, 10% split, model's own tokenizer
# parallel : ZeRO-3, DP=8 (single node, MI300X x 8)
# bucket : DeepSpeed defaults (stage3_prefetch_bucket_size = 5e7)
# seq/bs : seq_length=1024, micro_batch=1
# steps : 100 measured + 10 warmup
#
# Override via env vars: SEQ_LEN, BATCH_SIZE, NUM_STEPS, WARMUP_STEPS,
# NUM_GPUS, MODEL, DS_CONFIG.
set -eu
# Reduce HIP allocator fragmentation — the 32B model has long-lived tensors
# that benefit from segment expansion under heavy ZeRO-3 churn.
export PYTORCH_HIP_ALLOC_CONF=expandable_segments:True
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export TORCH_NCCL_ENABLE_MONITORING=0 # quiets harmless TCPStore shutdown trace
SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)
deepspeed --num_gpus "${NUM_GPUS:-8}" "${SCRIPT_DIR}/train_qwen3_zero3.py" \
--model_name "${MODEL:-Qwen/Qwen3-32B}" \
--num_layers "${NUM_LAYERS:-0}" \
--seq_length "${SEQ_LEN:-1024}" \
--batch_size "${BATCH_SIZE:-1}" \
--num_steps "${NUM_STEPS:-100}" \
--warmup_steps "${WARMUP_STEPS:-10}" \
--ds_config "${DS_CONFIG:-${SCRIPT_DIR}/ds_config_zero3.json}"
examples/sdma_allgather/run_qwen3_sdma_on.sh
Executable
+27
View File
@@ -0,0 +1,27 @@
#!/bin/bash
# Qwen3-32B + DeepSpeed ZeRO-3 with the SDMA fast-path opted in.
#
# DS_SDMA_ALLGATHER=1 is the single opt-in switch. When set,
# deepspeed.comm's TorchBackend tries to bring up the mori SDMA backend
# at init time and routes WORLD-group all_gather_into_tensor through it.
# Mori's MORI_ENABLE_SDMA=1 is auto-exported on the user's behalf when
# DS_SDMA_ALLGATHER=1 is set, so users normally don't need to touch it.
# This script otherwise uses the same ds_config as run_qwen3_sdma_off.sh;
# the only difference is this env var.
set -eu
export DS_SDMA_ALLGATHER=1
export PYTORCH_HIP_ALLOC_CONF=expandable_segments:True
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export TORCH_NCCL_ENABLE_MONITORING=0
SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)
deepspeed --num_gpus "${NUM_GPUS:-8}" "${SCRIPT_DIR}/train_qwen3_zero3.py" \
--model_name "${MODEL:-Qwen/Qwen3-32B}" \
--num_layers "${NUM_LAYERS:-0}" \
--seq_length "${SEQ_LEN:-1024}" \
--batch_size "${BATCH_SIZE:-1}" \
--num_steps "${NUM_STEPS:-100}" \
--warmup_steps "${WARMUP_STEPS:-10}" \
--ds_config "${DS_CONFIG:-${SCRIPT_DIR}/ds_config_zero3.json}"
examples/sdma_allgather/test_sdma_allgather_zero3.py
+198
View File
@@ -0,0 +1,198 @@
#!/usr/bin/env python3
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
"""
Unit test for the transparent SDMA allgather path in deepspeed.comm.
After ``deepspeed.init_distributed()`` returns, ``dist.all_gather_into_tensor``
on the WORLD process group is transparently routed through
``mori_cpp.AllGatherIntoTensor`` on AMD MI300 when mori is available, with
RCCL/NCCL as a fallback. This test exercises that path the same way
ZeRO-3's ``_all_gather_dtype`` does (flat output / per-rank shard input
with ``async_op=True``) and verifies correctness and algorithm bandwidth
for the common dtypes.
Usage:
cd examples/sdma_allgather
deepspeed --num_gpus 8 test_sdma_allgather_zero3.py
deepspeed --num_gpus 8 test_sdma_allgather_zero3.py --partition_sz 4194304 --iterations 50
"""
import argparse
import os
import numpy as np
import torch
import deepspeed
from deepspeed import comm as dist
from deepspeed.accelerator import get_accelerator
from deepspeed.comm import mori as _mori
def verify_allgather(partitions, world_size, partition_sz, rank, dtype):
"""Verify that each rank's partition contains the expected fill pattern."""
passed = True
for r in range(world_size):
chunk = partitions[r].narrow(0, 0, partition_sz).float().cpu()
expected_val = float(r + 1)
if not torch.allclose(chunk, torch.full_like(chunk, expected_val)):
unique_vals = chunk.unique()
print(f" [rank {rank}] FAIL: partition[{r}] expected all {expected_val}, "
f"got unique values: {unique_vals[:10]}")
passed = False
return passed
def run_single_allgather(rank, world_size, dtype, partition_sz, ag_stream):
"""Execute one allgather call following the ZeRO-3 ``_all_gather_dtype`` path."""
device = get_accelerator().current_device_name()
flat_tensor = torch.empty(partition_sz * world_size, dtype=dtype, device=device, requires_grad=False)
partitions = [flat_tensor.narrow(0, partition_sz * i, partition_sz) for i in range(world_size)]
partitions[rank].fill_(float(rank + 1))
with get_accelerator().stream(ag_stream):
handle = dist.allgather_fn(flat_tensor, partitions[rank], async_op=True)
with get_accelerator().stream(ag_stream):
handle.wait()
get_accelerator().current_stream().wait_stream(ag_stream)
return partitions
def run_bandwidth_test(rank, world_size, dtype, partition_sz, ag_stream, iterations, warmup):
"""Measure allgather bandwidth following the ZeRO-3 overlap pattern."""
device = get_accelerator().current_device_name()
elem_size = torch.tensor([], dtype=dtype).element_size()
total_bytes = partition_sz * elem_size * world_size
ev_start = get_accelerator().Event(enable_timing=True)
ev_end = get_accelerator().Event(enable_timing=True)
times_ms = []
for i in range(warmup + iterations):
flat_tensor = torch.empty(partition_sz * world_size, dtype=dtype, device=device, requires_grad=False)
partitions = [flat_tensor.narrow(0, partition_sz * r, partition_sz) for r in range(world_size)]
partitions[rank].fill_(float(rank + 1))
dist.barrier()
ev_start.record(ag_stream)
with get_accelerator().stream(ag_stream):
handle = dist.allgather_fn(flat_tensor, partitions[rank], async_op=True)
with get_accelerator().stream(ag_stream):
handle.wait()
ev_end.record(ag_stream)
ag_stream.synchronize()
elapsed_ms = ev_start.elapsed_time(ev_end)
if i >= warmup:
times_ms.append(elapsed_ms)
return times_ms, total_bytes
def main():
parser = argparse.ArgumentParser(description="Transparent SDMA allgather unit test")
parser.add_argument("--partition_sz", type=int, default=1024 * 1024, help="Elements per rank per allgather call")
parser.add_argument("--iterations", type=int, default=20, help="Number of measurement iterations")
parser.add_argument("--warmup", type=int, default=5, help="Number of warmup iterations")
parser.add_argument("--local_rank", type=int, default=int(os.environ.get("LOCAL_RANK", 0)))
parser = deepspeed.add_config_arguments(parser)
args = parser.parse_args()
deepspeed.init_distributed(dist_backend="cpu:gloo,cuda:nccl")
rank = dist.get_rank()
world_size = dist.get_world_size()
get_accelerator().set_device(args.local_rank)
if rank == 0:
backend = "SDMA (mori)" if _mori.is_enabled() else "RCCL/NCCL (mori unavailable or disabled)"
print(f"\n{'=' * 65}")
print(f" Transparent SDMA Allgather Unit Test")
print(f" world_size : {world_size}")
print(f" partition_sz : {args.partition_sz:,} elements")
print(f" iterations : {args.iterations} (warmup {args.warmup})")
print(f" backend : {backend}")
print(f"{'=' * 65}\n")
ag_stream = get_accelerator().Stream()
test_dtypes = [
("bfloat16", torch.bfloat16),
("float16", torch.float16),
("float32", torch.float32),
]
if rank == 0:
print("--- Correctness ---")
all_correct = True
for dtype_name, dtype in test_dtypes:
dist.barrier()
partitions = run_single_allgather(rank, world_size, dtype, args.partition_sz, ag_stream)
passed = verify_allgather(partitions, world_size, args.partition_sz, rank, dtype)
passed_t = torch.tensor([1 if passed else 0], dtype=torch.int32)
dist.all_reduce(passed_t, op=dist.ReduceOp.MIN)
ok = passed_t.item() == 1
if rank == 0:
elem_bytes = torch.tensor([], dtype=dtype).element_size()
data_mb = args.partition_sz * elem_bytes * world_size / (1024**2)
status = "PASSED" if ok else "FAILED"
print(f" {dtype_name:10s} data={data_mb:8.2f} MB {status}")
if not ok:
all_correct = False
if rank == 0:
print(f"\n--- Bandwidth (iterations={args.iterations}, warmup={args.warmup}) ---")
print(f" {'dtype':10s} {'data_MB':>10s} {'avg_ms':>9s} "
f"{'min_ms':>9s} {'max_ms':>9s} {'algo_BW':>12s}")
print(f" {'-'*10} {'-'*10} {'-'*9} {'-'*9} {'-'*9} {'-'*12}")
for dtype_name, dtype in test_dtypes:
dist.barrier()
times_ms, total_bytes = run_bandwidth_test(rank, world_size, dtype, args.partition_sz, ag_stream,
args.iterations, args.warmup)
avg_ms = np.mean(times_ms)
min_ms = np.min(times_ms)
max_ms = np.max(times_ms)
avg_t = torch.tensor([avg_ms], dtype=torch.float64)
min_t = torch.tensor([min_ms], dtype=torch.float64)
max_t = torch.tensor([max_ms], dtype=torch.float64)
dist.all_reduce(avg_t, op=dist.ReduceOp.SUM)
dist.all_reduce(min_t, op=dist.ReduceOp.MIN)
dist.all_reduce(max_t, op=dist.ReduceOp.MAX)
if rank == 0:
g_avg_ms = avg_t.item() / world_size
g_min_ms = min_t.item()
g_max_ms = max_t.item()
data_mb = total_bytes / (1024**2)
algo_bw_gbs = total_bytes / (g_avg_ms / 1000) / (1024**3)
print(f" {dtype_name:10s} {data_mb:10.2f} {g_avg_ms:9.3f} "
f"{g_min_ms:9.3f} {g_max_ms:9.3f} {algo_bw_gbs:9.2f} GB/s")
dist.barrier()
if rank == 0:
print()
print(f"Result: {'All correctness tests PASSED' if all_correct else 'Some correctness tests FAILED'}")
print(f"{'=' * 65}\n")
get_accelerator().synchronize()
dist.barrier()
if _mori.is_enabled():
import mori.shmem as shmem
shmem.shmem_finalize()
if __name__ == "__main__":
main()
examples/sdma_allgather/train_qwen3_zero3.py
+283
View File
@@ -0,0 +1,283 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
"""Qwen3 + DeepSpeed ZeRO-3 benchmark for the SDMA allgather feature.
Loads a Qwen3 model with random initialisation under `deepspeed.zero.Init`
so each rank only allocates its 1/world_size shard, then runs a small number
of training steps on either real wikitext or synthetic random tokens. Step
time is measured rank-0 side and reported with peak memory and the average
loss. The same trainer is used for the SDMA-on and SDMA-off comparison runs
in run_qwen3_sdma_{on,off}.sh.
The SDMA fast-path is opt-in via a single env var: ``deepspeed.comm``
brings up the mori SDMA backend at init time when ``DS_SDMA_ALLGATHER=1``
and routes WORLD-group ``all_gather_into_tensor`` through
``mori_cpp.AllGatherIntoTensor`` on AMD MI300. No ``ds_config`` flag is
required. Leaving ``DS_SDMA_ALLGATHER`` unset (the default) reproduces
the RCCL/NCCL baseline for A/B comparisons even with mori installed.
Real-data path uses HuggingFace `datasets` to stream wikitext-103 and the
model's own tokenizer to pad/truncate to seq_length. No external benchmark
repo is required.
"""
import argparse
import os
import time
import deepspeed
import torch
from deepspeed import comm as dist
from deepspeed.accelerator import get_accelerator
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--model_name", default="Qwen/Qwen3-32B")
p.add_argument("--num_layers",
type=int,
default=0,
help="0 = use model default; smaller values for quick smoke runs")
p.add_argument("--seq_length", type=int, default=1024)
p.add_argument("--batch_size", type=int, default=1)
p.add_argument("--num_steps", type=int, default=50)
p.add_argument("--warmup_steps", type=int, default=10)
p.add_argument("--log_interval", type=int, default=10)
p.add_argument("--ds_config", required=True)
p.add_argument("--dataset",
default="wikitext",
choices=["wikitext", "synthetic"],
help="Real text (wikitext-103) or pre-generated random ids")
p.add_argument("--dataset_percentage",
type=float,
default=10.0,
help="Percentage of wikitext train split to load (1.0-100.0)")
p.add_argument("--local_rank", type=int, default=-1)
return p.parse_args()
# ---------------------------------------------------------------------------
# Self-contained data pipeline (no external benchmark repo dependency).
# ---------------------------------------------------------------------------
class _SyntheticDataset(Dataset):
"""Pre-generated random token ids for deterministic timing runs."""
def __init__(self, vocab_size, seq_length, num_samples=10000, seed=42):
gen = torch.Generator().manual_seed(seed)
self.input_ids = torch.randint(0, vocab_size, (num_samples, seq_length), generator=gen, dtype=torch.long)
self.seq_length = seq_length
def __len__(self):
return self.input_ids.shape[0]
def __getitem__(self, idx):
ids = self.input_ids[idx]
return {
"input_ids": ids,
"labels": ids.clone(),
"attention_mask": torch.ones(self.seq_length, dtype=torch.long),
}
def _build_wikitext_loader(model_name, seq_length, batch_size, dataset_percentage, rank, world_size, is_main):
"""Stream wikitext-103-raw-v1 as a concatenated token stream sliced into
fixed `seq_length` chunks.
This is the standard "group_texts" / GPT-style chunking pattern: every
sample is exactly seq_length REAL tokens with no padding and no per-row
boundaries. Result is uniform-difficulty samples, so the per-step loss
has no variance from "this row was 90 % padding" effects — which is what
made the per-row+padding variant of this loader jittery on bs=1 demos.
"""
from datasets import DownloadConfig, load_dataset
from datasets.utils.logging import disable_progress_bar
if not is_main:
disable_progress_bar()
fraction = max(1, int(dataset_percentage))
split = "train" if dataset_percentage >= 100 else f"train[:{fraction}%]"
if is_main:
print(f"[trainer] loading wikitext-103-raw-v1 split={split}")
raw = load_dataset("wikitext",
"wikitext-103-raw-v1",
split=split,
download_config=DownloadConfig(disable_tqdm=True))
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token or tokenizer.convert_ids_to_tokens(2)
if is_main:
print(f"[trainer] encoding {len(raw)} rows as a single stream ...")
text = "\n\n".join(t for t in raw["text"] if t.strip())
all_ids = tokenizer.encode(text, add_special_tokens=False)
# Optional cap on number of chunks (env var) so the per-epoch length can
# be tuned for short demos. 0 = use all available chunks.
max_chunks = int(os.environ.get("QWEN3_MAX_CHUNKS", "0"))
n_full = len(all_ids) // seq_length
if max_chunks > 0:
n_full = min(n_full, max_chunks)
chunks = torch.tensor(all_ids[:n_full * seq_length], dtype=torch.long).view(n_full, seq_length)
if is_main:
print(f"[trainer] chunked: {len(all_ids)} tokens -> {n_full} "
f"sequences of {seq_length} (no padding)",
flush=True)
class _ChunkDataset(Dataset):
def __init__(self, t):
self.t = t
def __len__(self):
return self.t.shape[0]
def __getitem__(self, idx):
ids = self.t[idx]
return {
"input_ids": ids,
"labels": ids.clone(),
"attention_mask": torch.ones(ids.shape[0], dtype=torch.long),
}
ds = _ChunkDataset(chunks)
sampler = DistributedSampler(ds, num_replicas=world_size, rank=rank)
return DataLoader(ds, batch_size=batch_size, sampler=sampler, num_workers=0, drop_last=True, pin_memory=True)
def _build_loader(args, vocab_size, rank, world_size, is_main):
if args.dataset == "wikitext":
return _build_wikitext_loader(args.model_name, args.seq_length, args.batch_size, args.dataset_percentage, rank,
world_size, is_main)
ds = _SyntheticDataset(vocab_size, args.seq_length)
return DataLoader(ds, batch_size=args.batch_size, shuffle=False, drop_last=True, num_workers=0, pin_memory=True)
# ---------------------------------------------------------------------------
# Model construction under deepspeed.zero.Init so each rank only materialises
# its shard.
# ---------------------------------------------------------------------------
def build_model(model_name, num_layers, ds_config_path):
cfg = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
if num_layers > 0:
cfg.num_hidden_layers = num_layers
cfg.torch_dtype = torch.bfloat16
cfg.use_cache = False
cfg.attn_implementation = "eager" # FA2 not always available on AMD; eager is safe.
if dist.is_initialized() and dist.get_rank() == 0:
print(f"[trainer] {model_name}: layers={cfg.num_hidden_layers} "
f"hidden={cfg.hidden_size} heads={cfg.num_attention_heads} "
f"kv_heads={cfg.num_key_value_heads} vocab={cfg.vocab_size}")
with deepspeed.zero.Init(config_dict_or_path=ds_config_path):
model = AutoModelForCausalLM.from_config(cfg, trust_remote_code=True)
return model, cfg
def main():
args = parse_args()
deepspeed.init_distributed()
rank = dist.get_rank()
world = dist.get_world_size()
accel = get_accelerator()
device_idx = args.local_rank if args.local_rank >= 0 else rank % accel.device_count()
device = torch.device(accel.device_name(device_idx))
accel.set_device(device_idx)
if rank == 0:
print(f"[trainer] world={world} device={device} ds_config={args.ds_config}")
model, cfg = build_model(args.model_name, args.num_layers, args.ds_config)
engine, _, _, _ = deepspeed.initialize(
args=args,
model=model,
model_parameters=[p for p in model.parameters() if p.requires_grad],
config=args.ds_config,
)
if rank == 0:
from deepspeed.comm import mori as _mori
print(f"[trainer] SDMA handle is_enabled={_mori.is_enabled()}", flush=True)
loader = _build_loader(args, cfg.vocab_size, rank, world, rank == 0)
if rank == 0:
print(f"[trainer] dataloader: {len(loader)} batches/epoch, "
f"running {args.num_steps} steps", flush=True)
step_times, losses = [], []
get_accelerator().reset_peak_memory_stats()
t_train_start = time.perf_counter()
step, epoch = 0, 0
data_iter = iter(loader)
skipped_empty = 0
while step < args.num_steps:
try:
batch = next(data_iter)
except StopIteration:
epoch += 1
if hasattr(loader.sampler, "set_epoch"):
loader.sampler.set_epoch(epoch)
data_iter = iter(loader)
batch = next(data_iter)
ids = batch["input_ids"].to(device, non_blocking=True)
labels = batch["labels"].to(device, non_blocking=True)
attn = batch["attention_mask"].to(device, non_blocking=True)
# Defensive: on the chunked wikitext loader every chunk is full of
# real tokens so these guards are no-ops, but they keep the trainer
# safe for the synthetic mode and any future padded variants.
if int(attn.sum().item()) == 0:
skipped_empty += 1
continue
labels = labels.masked_fill(attn == 0, -100)
get_accelerator().synchronize()
t0 = time.perf_counter()
out = engine(input_ids=ids, labels=labels, attention_mask=attn)
engine.backward(out.loss)
engine.step()
get_accelerator().synchronize()
dt = time.perf_counter() - t0
if step >= args.warmup_steps:
step_times.append(dt)
losses.append(out.loss.detach().item())
if rank == 0 and step % args.log_interval == 0:
tag = "warmup" if step < args.warmup_steps else "measured"
tps = args.batch_size * args.seq_length * world / dt
print(
f"[trainer] step {step:4d} ({tag:7s}) | loss {out.loss.item():8.4f} | "
f"step {dt*1000:7.1f} ms | {tps:8.0f} tok/s",
flush=True)
step += 1
t_train_end = time.perf_counter()
if rank == 0:
n = len(step_times)
avg_dt = sum(step_times) / n
tokens_per_step = args.batch_size * args.seq_length * world
tps = tokens_per_step / avg_dt
peak_gb = get_accelerator().max_memory_allocated() / 1e9
avg_loss = sum(losses) / n
print()
print("=" * 70)
print("Qwen3 ZeRO-3 benchmark complete")
print(f" measured steps : {n} (warmup={args.warmup_steps} skipped)")
print(f" total wall (s) : {t_train_end - t_train_start:.1f}")
print(f" avg step (ms) : {avg_dt * 1000:.1f}")
print(f" tokens/sec (ws) : {tps:.1f}")
print(f" peak mem (GB,r0) : {peak_gb:.2f}")
print(f" avg loss : {avg_loss:.4f}")
print(f" final loss : {losses[-1]:.4f}")
print(f" empty batches : {skipped_empty}")
print("=" * 70)
if __name__ == "__main__":
main()
examples/sdma_allgather/train_zero3.py
+335
View File
@@ -0,0 +1,335 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
"""
DeepSpeed ZeRO-3 training example with allgather overlap.
Trains a GPT-2-style transformer on synthetic data for demonstration.
Designed for single-node 8x AMD GPU setup.
"""
import argparse
import math
import os
import time
import torch
import torch.nn as nn
import deepspeed
from deepspeed import comm as dist
from deepspeed.accelerator import get_accelerator
from torch.utils.data import Dataset, DataLoader
# ---------------------------------------------------------------------------
# Model: minimal GPT-2-style transformer
# ---------------------------------------------------------------------------
class CausalSelfAttention(nn.Module):
def __init__(self, hidden_size, num_heads, max_seq_len, dropout=0.1):
super().__init__()
assert hidden_size % num_heads == 0
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.qkv = nn.Linear(hidden_size, 3 * hidden_size)
self.proj = nn.Linear(hidden_size, hidden_size)
self.attn_drop = nn.Dropout(dropout)
self.proj_drop = nn.Dropout(dropout)
self.register_buffer(
"causal_mask",
torch.tril(torch.ones(max_seq_len, max_seq_len)).view(1, 1, max_seq_len, max_seq_len),
)
def forward(self, x):
B, T, C = x.size()
q, k, v = self.qkv(x).split(C, dim=-1)
q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
scale = 1.0 / math.sqrt(self.head_dim)
attn = (q @ k.transpose(-2, -1)) * scale
attn = attn.masked_fill(self.causal_mask[:, :, :T, :T] == 0, float("-inf"))
attn = torch.softmax(attn, dim=-1)
attn = self.attn_drop(attn)
out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C)
return self.proj_drop(self.proj(out))
class TransformerBlock(nn.Module):
def __init__(self, hidden_size, num_heads, max_seq_len, dropout=0.1):
super().__init__()
self.ln1 = nn.LayerNorm(hidden_size)
self.attn = CausalSelfAttention(hidden_size, num_heads, max_seq_len, dropout)
self.ln2 = nn.LayerNorm(hidden_size)
self.mlp = nn.Sequential(
nn.Linear(hidden_size, 4 * hidden_size),
nn.GELU(),
nn.Linear(4 * hidden_size, hidden_size),
nn.Dropout(dropout),
)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
class GPT2Model(nn.Module):
def __init__(self, vocab_size, hidden_size, num_layers, num_heads, max_seq_len, dropout=0.1):
super().__init__()
self.tok_emb = nn.Embedding(vocab_size, hidden_size)
self.pos_emb = nn.Embedding(max_seq_len, hidden_size)
self.drop = nn.Dropout(dropout)
self.blocks = nn.Sequential(
*[TransformerBlock(hidden_size, num_heads, max_seq_len, dropout) for _ in range(num_layers)])
self.ln_f = nn.LayerNorm(hidden_size)
self.head = nn.Linear(hidden_size, vocab_size, bias=False)
def forward(self, input_ids, labels=None):
B, T = input_ids.size()
pos = torch.arange(0, T, device=input_ids.device).unsqueeze(0)
x = self.drop(self.tok_emb(input_ids) + self.pos_emb(pos))
x = self.blocks(x)
x = self.ln_f(x)
logits = self.head(x)
loss = None
if labels is not None:
loss = nn.functional.cross_entropy(
logits.view(-1, logits.size(-1)),
labels.view(-1),
)
return loss, logits
# ---------------------------------------------------------------------------
# Synthetic dataset
# ---------------------------------------------------------------------------
class SyntheticTextDataset(Dataset):
"""Generates synthetic token sequences for perf/correctness testing."""
def __init__(self, vocab_size, seq_len, num_samples, seed=42, mode="random"):
self.vocab_size = vocab_size
self.seq_len = seq_len
self.num_samples = num_samples
self.seed = seed
self.mode = mode
def __len__(self):
return self.num_samples
def __getitem__(self, idx):
if self.mode == "random":
g = torch.Generator()
g.manual_seed(self.seed + idx)
tokens = torch.randint(0, self.vocab_size, (self.seq_len + 1, ), generator=g)
elif self.mode == "arange":
start = (self.seed + idx) % self.vocab_size
tokens = (torch.arange(self.seq_len + 1, dtype=torch.long) + start) % self.vocab_size
elif self.mode == "repeat":
v = (self.seed + idx) % self.vocab_size
tokens = torch.full((self.seq_len + 1, ), v, dtype=torch.long)
else:
raise ValueError(f"Unsupported data mode: {self.mode}")
return tokens[:-1], tokens[1:]
class WikitextDataset(Dataset):
"""Real text dataset from HuggingFace wikitext-2 / wikitext-103."""
def __init__(self, vocab_size, seq_len, num_samples, split="train", dataset_name="wikitext-2-raw-v1"):
from datasets import load_dataset
from transformers import GPT2TokenizerFast
tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
raw = load_dataset("wikitext", dataset_name, split=split)
text = "\n\n".join([t for t in raw["text"] if t.strip()])
all_ids = tokenizer.encode(text)
self.seq_len = seq_len
self.samples = []
for i in range(0, len(all_ids) - seq_len - 1, seq_len):
self.samples.append(torch.tensor(all_ids[i:i + seq_len + 1], dtype=torch.long))
if len(self.samples) >= num_samples:
break
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
tokens = self.samples[idx]
return tokens[:-1], tokens[1:]
# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def parse_args():
parser = argparse.ArgumentParser(description="DeepSpeed ZeRO-3 training with allgather overlap")
parser.add_argument("--vocab_size", type=int, default=50257)
parser.add_argument("--hidden_size", type=int, default=4096)
parser.add_argument("--num_layers", type=int, default=48)
parser.add_argument("--num_heads", type=int, default=32)
parser.add_argument("--max_seq_len", type=int, default=2048)
parser.add_argument("--dropout", type=float, default=0.1)
parser.add_argument("--num_samples", type=int, default=10000)
parser.add_argument("--train_steps", type=int, default=2000)
parser.add_argument("--data_mode",
type=str,
default="random",
choices=["random", "arange", "repeat", "wikitext2", "wikitext103"],
help="Data mode. random/arange/repeat are synthetic; wikitext2/wikitext103 use real text.")
parser.add_argument("--local_rank", type=int, default=-1)
parser = deepspeed.add_config_arguments(parser)
return parser.parse_args()
def main():
args = parse_args()
ds_config_path = args.deepspeed_config
if ds_config_path and not os.path.isfile(ds_config_path):
script_dir = os.path.dirname(os.path.abspath(__file__))
ds_config_path = os.path.join(script_dir, ds_config_path)
args.deepspeed_config = ds_config_path
deepspeed.init_distributed(dist_backend="cpu:gloo,cuda:nccl")
local_rank = args.local_rank
get_accelerator().set_device(local_rank)
torch.manual_seed(42)
get_accelerator().manual_seed_all(42)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
with deepspeed.zero.Init(config_dict_or_path=ds_config_path):
model = GPT2Model(
vocab_size=args.vocab_size,
hidden_size=args.hidden_size,
num_layers=args.num_layers,
num_heads=args.num_heads,
max_seq_len=args.max_seq_len,
dropout=0.0,
)
total_params = sum(p.numel() for p in model.parameters())
num_gpus = dist.get_world_size()
if local_rank == 0:
print(f"Model parameters: {total_params / 1e6:.1f}M")
print(f"GPUs: {num_gpus}")
# FLOPs per token (forward + backward): 6*params + 12*L*H*S
# Reference: "Efficient Large-Scale Language Model Training on GPU Clusters
# Using Megatron-LM" (Narayanan et al., 2021)
flops_per_token = 6 * total_params + 12 * args.num_layers * args.hidden_size * args.max_seq_len
if args.data_mode in ("wikitext2", "wikitext103"):
ds_name = "wikitext-2-raw-v1" if args.data_mode == "wikitext2" else "wikitext-103-raw-v1"
dataset = WikitextDataset(args.vocab_size, args.max_seq_len, args.num_samples, dataset_name=ds_name)
else:
dataset = SyntheticTextDataset(args.vocab_size, args.max_seq_len, args.num_samples, mode=args.data_mode)
if local_rank == 0:
if args.data_mode == "random":
print(f"Data mode: random (expected CE floor ~ ln(vocab) = {math.log(args.vocab_size):.4f})")
elif args.data_mode in ("wikitext2", "wikitext103"):
print(f"Data mode: {args.data_mode} (real text, {len(dataset)} samples)")
else:
print(f"Data mode: {args.data_mode} (learnable pattern, loss should decrease)")
model_engine, optimizer, _, lr_scheduler = deepspeed.initialize(
args=args,
model=model,
)
sampler = torch.utils.data.distributed.DistributedSampler(
dataset,
shuffle=False,
seed=42,
)
train_loader = DataLoader(
dataset,
batch_size=model_engine.train_micro_batch_size_per_gpu(),
sampler=sampler,
num_workers=0,
pin_memory=True,
)
device = model_engine.device
global_batch = model_engine.train_batch_size()
tokens_per_step = global_batch * args.max_seq_len
warmup_steps = min(50, args.train_steps // 10)
step = 0
step_times = []
t_start = time.time()
t_steady = None
while step < args.train_steps:
for batch in train_loader:
if step >= args.train_steps:
break
get_accelerator().synchronize()
t_step_start = time.time()
input_ids = batch[0].to(device)
labels = batch[1].to(device)
loss, _ = model_engine(input_ids, labels=labels)
model_engine.backward(loss)
model_engine.step()
get_accelerator().synchronize()
step_time_ms = (time.time() - t_step_start) * 1000
if step == warmup_steps:
t_steady = time.time()
if step >= warmup_steps:
step_times.append(step_time_ms)
if step % 10 == 0 and local_rank == 0:
if step_times:
import numpy as np
recent = np.array(step_times[-20:])
avg_ms = recent.mean()
cur_samples_per_sec = global_batch / (avg_ms / 1000)
cur_tokens_per_sec = cur_samples_per_sec * args.max_seq_len
cur_tflops_per_gpu = cur_tokens_per_sec * flops_per_token / 1e12 / num_gpus
else:
avg_ms = step_time_ms
cur_tflops_per_gpu = 0.0
cur_samples_per_sec = 0.0
print(f"step {step:5d} | loss {loss.item():.4f} | "
f"lr {lr_scheduler.get_last_lr()[0]:.6f} | "
f"{cur_samples_per_sec:.1f} samples/s | "
f"{cur_tflops_per_gpu:.2f} TFLOPS/GPU | "
f"step {avg_ms:.1f} ms")
step += 1
if local_rank == 0:
import numpy as np
total_time = time.time() - t_start
st = np.array(step_times)
steady_steps = len(st)
steady_time = time.time() - t_steady if t_steady else total_time
steady_samples_per_sec = steady_steps * global_batch / steady_time
steady_tokens_per_sec = steady_samples_per_sec * args.max_seq_len
steady_tflops = steady_tokens_per_sec * flops_per_token / 1e12
steady_tflops_per_gpu = steady_tflops / num_gpus
print(f"\n{'=' * 70}")
print(f"Training complete: {args.train_steps} steps in {total_time:.1f}s")
print(f" (warmup={warmup_steps} steps skipped, measured {steady_steps} steps)")
print(f"{'=' * 70}")
print(f" Throughput : {steady_samples_per_sec:.1f} samples/s")
print(f" TFLOPS : {steady_tflops:.1f} (total) | {steady_tflops_per_gpu:.2f} (per GPU)")
print(f" Step time (ms) : avg {st.mean():.1f} | p50 {np.median(st):.1f} | "
f"p99 {np.percentile(st, 99):.1f} | min {st.min():.1f} | max {st.max():.1f}")
print(f"{'=' * 70}")
if __name__ == "__main__":
main()