wehub-resource-sync
229 lines
8.4 KiB
Python
229 lines
8.4 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""mori SDMA backend, plugged into ``TorchBackend.all_gather_into_tensor``.
|
|
|
|
When the user opts in, ``deepspeed.comm`` routes ``all_gather_into_tensor``
|
|
on the WORLD process group through ``mori_cpp.AllGatherIntoTensor``
|
|
(intra-node SDMA copy on AMD MI300). Any failure (mori missing,
|
|
non-AMD/ROCm runtime, shmem init error, oversized call, non-WORLD group)
|
|
yields ``None`` and the caller falls back to the underlying RCCL/NCCL
|
|
allgather.
|
|
|
|
User-visible controls (env vars, no ``ds_config`` field):
|
|
|
|
* ``DS_SDMA_ALLGATHER=1`` opt in to the SDMA path. Required:
|
|
even when mori is installed, the
|
|
SDMA fast-path stays off unless
|
|
the user sets this explicitly.
|
|
When set, ``MORI_ENABLE_SDMA=1`` is
|
|
auto-exported on the user's behalf
|
|
so mori allocates uncached transit
|
|
buffers.
|
|
* ``DS_SDMA_ALLGATHER_MAX_NUMEL=N`` override the transit buffer size in
|
|
elements (default 64M = 256 MiB
|
|
per-rank input, ~2 GiB output on 8
|
|
ranks)
|
|
"""
|
|
|
|
import os
|
|
from typing import Optional
|
|
|
|
import torch
|
|
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.utils import logger
|
|
|
|
_handle = None
|
|
_dtype_map = None
|
|
_max_numel = 0
|
|
_init_attempted = False
|
|
_call_failed_warned = False
|
|
|
|
|
|
class _SdmaWork:
|
|
"""Duck-type compatible with ``torch.distributed.Work``.
|
|
|
|
``wait()`` issues a stream-level dependency only and does NOT block the
|
|
CPU, mirroring RCCL ``Work.wait()`` semantics. ZeRO-3's prefetch
|
|
pipeline relies on the CPU staying free so the next bucket can be
|
|
queued ahead of time while bucket N is in flight.
|
|
"""
|
|
|
|
def __init__(self, event):
|
|
self._event = event
|
|
|
|
def wait(self):
|
|
get_accelerator().current_stream().wait_event(self._event)
|
|
|
|
def is_completed(self) -> bool:
|
|
return self._event.query()
|
|
|
|
|
|
def _ensure_default_pg_registered():
|
|
"""Register the WORLD process group as 'default' in PyTorch's C++ GroupRegistry.
|
|
|
|
mori's shmem layer looks up the PG by the name "default"; the standard
|
|
DeepSpeed init path doesn't register WORLD under that label.
|
|
"""
|
|
world_group = torch.distributed.group.WORLD
|
|
assert world_group is not None, "torch.distributed must be initialized before SDMA allgather"
|
|
torch._C._distributed_c10d._register_process_group("default", world_group)
|
|
|
|
|
|
def _build_dtype_map():
|
|
"""torch.dtype -> mori_cpp.DataType (NCCL-style enum)."""
|
|
from mori.ccl import DataType
|
|
return {
|
|
torch.uint8: DataType.Uint8,
|
|
torch.int8: DataType.Int8,
|
|
torch.int16: DataType.Int16,
|
|
torch.int32: DataType.Int32,
|
|
torch.int64: DataType.Int64,
|
|
torch.float16: DataType.Float16,
|
|
torch.bfloat16: DataType.BFloat16,
|
|
torch.float32: DataType.Float32,
|
|
torch.float64: DataType.Float64,
|
|
}
|
|
|
|
|
|
_TRUTHY = {"1", "true", "True", "TRUE", "yes", "Yes", "YES", "on", "On", "ON"}
|
|
|
|
|
|
def _is_enabled_by_env() -> bool:
|
|
"""User must explicitly opt in via ``DS_SDMA_ALLGATHER=1``.
|
|
|
|
Default is off even when mori happens to be importable: mori is an
|
|
external dependency and we don't want DeepSpeed's collective backend
|
|
to silently change behaviour based on which extra packages are
|
|
installed. Keeping this opt-in also makes A/B baselines against the
|
|
stock RCCL path trivial without having to uninstall mori.
|
|
"""
|
|
return os.environ.get("DS_SDMA_ALLGATHER", "0") in _TRUTHY
|
|
|
|
|
|
def _resolve_max_numel(default: int) -> int:
|
|
raw = os.environ.get("DS_SDMA_ALLGATHER_MAX_NUMEL")
|
|
if raw is None:
|
|
return default
|
|
try:
|
|
return max(int(raw), 0)
|
|
except ValueError:
|
|
return default
|
|
|
|
|
|
def init(max_numel: int = 64 * 1024 * 1024) -> None:
|
|
"""Best-effort, idempotent SDMA handle construction.
|
|
|
|
Builds one ``mori_cpp.AllGatherIntoTensor`` (NCCL/RCCL-style C++
|
|
dispatcher) sized for the largest expected per-rank shard. All
|
|
subsequent allgather calls reuse this handle. Safe to call
|
|
unconditionally: any failure leaves ``_handle`` unset and logs a
|
|
single rank-0 info line, so callers transparently fall back to
|
|
RCCL/NCCL.
|
|
"""
|
|
global _handle, _dtype_map, _max_numel, _init_attempted
|
|
if _init_attempted:
|
|
return
|
|
_init_attempted = True
|
|
|
|
is_rank0 = torch.distributed.is_initialized() and torch.distributed.get_rank() == 0
|
|
if not _is_enabled_by_env():
|
|
# Silent no-op: SDMA stays off and dist.allgather is used. We
|
|
# don't log here because most users never set DS_SDMA_ALLGATHER and
|
|
# rank-0 spam on every backend init is noise.
|
|
return
|
|
|
|
max_numel = _resolve_max_numel(max_numel)
|
|
# mori's SymmMemManager only allocates the uncached transit buffers
|
|
# required by the SDMA kernel when MORI_ENABLE_SDMA is set; setdefault
|
|
# so users who already exported it (or want to override) win.
|
|
os.environ.setdefault("MORI_ENABLE_SDMA", "1")
|
|
|
|
try:
|
|
_ensure_default_pg_registered()
|
|
import mori.shmem as shmem
|
|
from mori.ccl import AllGatherIntoTensor
|
|
|
|
shmem.shmem_torch_process_group_init("default")
|
|
my_pe = shmem.shmem_mype()
|
|
npes = shmem.shmem_npes()
|
|
# Per-rank input transit buffer must hold the largest shard we'll
|
|
# ever see; output buffer = npes * input. 4 B/element is the SDMA
|
|
# kernel's uint32 lane width.
|
|
input_bytes = max_numel * 4
|
|
_handle = AllGatherIntoTensor(
|
|
my_pe=my_pe,
|
|
npes=npes,
|
|
input_buffer_size=input_bytes,
|
|
output_buffer_size=input_bytes * npes,
|
|
copy_output_to_user=True,
|
|
)
|
|
_dtype_map = _build_dtype_map()
|
|
_max_numel = max_numel
|
|
if is_rank0:
|
|
logger.info(f"SDMA allgather enabled via mori_cpp.AllGatherIntoTensor "
|
|
f"(max_numel={max_numel})")
|
|
except Exception as e:
|
|
_handle = None
|
|
_dtype_map = None
|
|
_max_numel = 0
|
|
if is_rank0:
|
|
logger.info(f"SDMA allgather unavailable ({type(e).__name__}: {e}); "
|
|
f"using RCCL/NCCL allgather")
|
|
|
|
|
|
def is_enabled() -> bool:
|
|
return _handle is not None
|
|
|
|
|
|
def supports(input_tensor: torch.Tensor, group=None) -> bool:
|
|
"""Cheap pre-check used by ``TorchBackend.all_gather_into_tensor``.
|
|
|
|
SDMA is only safe when:
|
|
- the backend is initialised (``_handle`` set),
|
|
- the call is on the WORLD process group (mori's shmem layer was
|
|
bound to "default"/WORLD at init time),
|
|
- the per-rank shard fits inside the pre-allocated transit buffer,
|
|
- the input dtype is in ``_dtype_map``.
|
|
"""
|
|
if _handle is None:
|
|
return False
|
|
if group is not None and group is not torch.distributed.group.WORLD:
|
|
return False
|
|
if input_tensor.numel() > _max_numel:
|
|
return False
|
|
if _dtype_map is None or input_tensor.dtype not in _dtype_map:
|
|
return False
|
|
return True
|
|
|
|
|
|
def allgather_into_tensor(input_tensor: torch.Tensor, output_tensor: torch.Tensor, group=None) -> Optional[_SdmaWork]:
|
|
"""Run one allgather_into_tensor through the SDMA handle.
|
|
|
|
Returns an ``_SdmaWork`` (Work-compatible) on success. Returns
|
|
``None`` when SDMA is not applicable for this call (uninitialised,
|
|
non-WORLD group, dtype not supported, shard larger than the transit
|
|
buffer) or the call fails for any reason — the caller falls back to
|
|
``dist.allgather_fn``.
|
|
"""
|
|
global _call_failed_warned
|
|
if not supports(input_tensor, group):
|
|
return None
|
|
try:
|
|
stream = get_accelerator().current_stream()
|
|
dtype = _dtype_map[input_tensor.dtype]
|
|
ok = _handle(input_tensor.data_ptr(), output_tensor.data_ptr(), input_tensor.numel(), dtype,
|
|
stream.cuda_stream)
|
|
if not ok:
|
|
return None
|
|
event = get_accelerator().Event()
|
|
event.record(stream)
|
|
return _SdmaWork(event)
|
|
except Exception as e:
|
|
if (not _call_failed_warned and torch.distributed.is_initialized() and torch.distributed.get_rank() == 0):
|
|
logger.warning(f"SDMA allgather failed ({e}); falling back to dist.allgather")
|
|
_call_failed_warned = True
|
|
return None
|