diff --git a/docker/deepseekv32/Dockerfile b/docker/deepseekv32/Dockerfile new file mode 100644 index 00000000000..38e9eadfafd --- /dev/null +++ b/docker/deepseekv32/Dockerfile @@ -0,0 +1,145 @@ +ARG SGLANG_IMAGE_TAG=v0.5.6.post2 +FROM lmsysorg/sglang:${SGLANG_IMAGE_TAG} AS sglang + +# ======================================== Arguments ============================================= + +ARG PATCH_VERSION=latest +ARG MEGATRON_COMMIT=436065a86b749ca3b50eebca68f55c9e690a9f63 + +ARG ENABLE_CUDA_13=0 + +ARG ENABLE_SGLANG_PATCH=0 + +# ======================================== Setup ============================================= + +WORKDIR /root/ + +# ======================================== Apt dependencies ============================================= + +RUN apt update +RUN apt install -y nvtop rsync dnsutils + +# ====================================== Python dependencies ============================================ + +# The compilation is slow, thus should be put at top +# TransformerEngines does not support too high FA2 +RUN MAX_JOBS=64 pip -v install flash-attn==2.7.4.post1 --no-build-isolation + +# The compilation is slow, thus should be put at top +RUN git clone https://github.com/Dao-AILab/flash-attention.git && \ + cd flash-attention/ && git checkout fbf24f67cf7f6442c5cfb2c1057f4bfc57e72d89 && git submodule update --init && cd hopper/ && \ + MAX_JOBS=96 python setup.py install && \ + export python_path=`python -c "import site; print(site.getsitepackages()[0])"` && \ + mkdir -p $python_path/flash_attn_3 && \ + cp flash_attn_interface.py $python_path/flash_attn_3/flash_attn_interface.py && \ + rm -rf flash-attention/ + +RUN pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c --no-deps + +RUN pip install flash-linear-attention==0.4.0 + +RUN git clone https://github.com/Dao-AILab/fast-hadamard-transform.git fast-hadamard-transform && \ + cd fast-hadamard-transform && \ + pip install -v . --no-build-isolation && \ + cd /root && \ + rm -rf fast-hadamard-transform + +# TE does not have wheel on cuda 13 yet, thus need to install from source +RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \ + pip install nvidia-mathdx==25.6.0 && \ + pip install pybind11 && \ + pip -v install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.10; \ + else \ + pip -v install --no-build-isolation "transformer_engine[pytorch]==2.10.0"; \ + fi + +RUN NVCC_APPEND_FLAGS="--threads 4" \ + pip -v install --disable-pip-version-check --no-cache-dir \ + --no-build-isolation \ + --config-settings "--build-option=--cpp_ext --cuda_ext --parallel 8" git+https://github.com/NVIDIA/apex.git@10417aceddd7d5d05d7cbf7b0fc2daad1105f8b4 + +RUN git clone https://github.com/NVIDIA/Megatron-LM.git --recursive && \ + cd Megatron-LM && git checkout ${MEGATRON_COMMIT} && \ + pip install -e . + +RUN git clone https://github.com/huggingface/transformers.git && \ + cd transformers && git checkout 8cb5963cc22174954e7dca2c0a3320b7dc2f4edc && \ + pip install -e . + +RUN pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@dc6876905830430b5054325fa4211ff302169c6b --no-cache-dir --force-reinstall +RUN pip install git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl --no-build-isolation +RUN pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation + +# This patch from masahi will be included in later Triton releases +RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \ + (cd /root && git clone -b feat/v350_plus_8045 https://github.com/fzyzcjy/triton.git && cd triton && pip install -r python/requirements.txt && pip install --verbose -e .); \ + fi + +COPY requirements.txt /tmp/requirements.txt +RUN pip install -r /tmp/requirements.txt + +# Temporarily install another sgl-kernel version for GB300 without rebuilding the whole image +RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \ + SGL_KERNEL_VERSION=0.3.17.post2 && \ + python3 -m pip install https://github.com/sgl-project/whl/releases/download/v${SGL_KERNEL_VERSION}/sgl_kernel-${SGL_KERNEL_VERSION}+cu130-cp310-abi3-manylinux2014_$(uname -m).whl --force-reinstall --no-deps; \ + fi + +# AMEM +# we need to create a fake libcuda.so.1 to make the linker happy when building AMEM +# ENV CUDA_DIR=/usr/local/cuda +# ENV CUDA_STUBS=${CUDA_DIR}/lib64/stubs +# RUN ln -s ${CUDA_STUBS}/libcuda.so ${CUDA_STUBS}/libcuda.so.1 && \ +# echo "${CUDA_STUBS}" > /etc/ld.so.conf.d/z-cuda-stubs.conf && \ +# ldconfig +# RUN git clone https://github.com/inclusionAI/asystem-amem.git && \ +# cd asystem-amem && git checkout 6483bb17c9a98b51c3a94b7048467d5b50fbad4b && \ +# git submodule init && git submodule update && \ +# MPI_HOME=/usr/lib/x86_64-linux-gnu/openmpi/ ./build.sh && \ +# mv /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/libnccl.so.2 /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/libnccl.so.2.bak && \ +# cp -r third_party/nccl/build/lib/* /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/ + +RUN [ ! -f /root/.tmux.conf ] || rm /root/.tmux.conf + +# ====================================== Patches ============================================ + +COPY docker/deepseekv32/megatron.patch /root/Megatron-LM/ +RUN cd Megatron-LM && \ + git update-index --refresh && \ + git apply megatron.patch --3way && \ + if grep -R -n '^<<<<<<< ' .; then \ + echo "Patch failed to apply cleanly. Please resolve conflicts." && \ + exit 1; \ + fi && \ + rm megatron.patch + +COPY docker/deepseekv32/transformers.patch /root/transformers/ +RUN cd transformers && \ + git update-index --refresh && \ + git apply transformers.patch --3way && \ + if grep -R -n '^<<<<<<< ' .; then \ + echo "Patch failed to apply cleanly. Please resolve conflicts." && \ + exit 1; \ + fi && \ + rm transformers.patch + +# TODO temporarily skip patching for GB200/GB300 (and require users to bring their own sglang version). should add back later. +COPY docker/patch/${PATCH_VERSION}/sglang.patch /sgl-workspace/sglang/ +RUN if [ "$ENABLE_SGLANG_PATCH" = "1" ]; then \ + cd /sgl-workspace/sglang && \ + git update-index --refresh && \ + git apply sglang.patch && \ + if grep -R -n '^<<<<<<< ' .; then \ + echo "Patch failed to apply cleanly. Please resolve conflicts." && \ + exit 1; \ + fi && \ + rm sglang.patch; \ +fi + +# ====================================== Install main package ============================================ + +# TODO may improve +ARG MILES_COMMIT=main +RUN git clone https://github.com/radixark/miles.git /root/miles && \ + cd /root/miles && \ + git checkout ${MILES_COMMIT} && \ + pip install -e . --no-deps diff --git a/docker/deepseekv32/README.md b/docker/deepseekv32/README.md new file mode 100644 index 00000000000..1ca43f1c2bd --- /dev/null +++ b/docker/deepseekv32/README.md @@ -0,0 +1,41 @@ +## Usage + +### Docker +```bash +docker pull yueming11/miles:dsv32-dev + +docker run --gpus all --ipc=host --shm-size=16g --ulimit memlock=-1 --ulimit stack=67108864 --name miles_dsv32 yueming11/miles:dsv32-dev /bin/zsh + +git clone https://github.com/radixark/miles.git +git checkout dsv32 +cd dsv32 +pip install -e . + +# if shows Megatron does not support numpy 2.x +pip install numpy==1.26.4 +``` + +### Quick test with 5 layer model +#### model download + +``` +hf download Pinaster/DeepSeek-V3.2-5layer /root/models/DeepSeek-V3.2-5layer +``` + +#### Prepare model for training +Note: need to change the paths, for all commands below see `scripts/run_deepseek_v32.py` for details + +Step 1. download dataset & convert fp8 hf checkpoint to bf16 with one node +``` +python scripts/run_deepseek_v32.py prepare-single --model-name DeepSeek-V3.2-5layer --megatron-model-type deepseek-v32-5layer +``` + +Step 2. convert hf checkpoint to megatron checkpoint with multiple nodes +``` +python scripts/run_deepseek_v32.py prepare-spmd --model-name DeepSeek-V3.2-5layer --megatron-model-type deepseek-v32-5layer +``` + +#### Launch training +``` +python scripts/run_deepseek_v32.py train --model-name DeepSeek-V3.2-5layer --megatron-model-type deepseek-v32-5layer +``` \ No newline at end of file diff --git a/docker/deepseekv32/megatron.patch b/docker/deepseekv32/megatron.patch new file mode 100644 index 00000000000..f7204197c3a --- /dev/null +++ b/docker/deepseekv32/megatron.patch @@ -0,0 +1,1559 @@ +diff --git a/megatron/core/transformer/dot_product_attention_context_parallel.py b/megatron/core/transformer/dot_product_attention_context_parallel.py +index 89659a1d7..c69859a04 100644 +--- a/megatron/core/transformer/dot_product_attention_context_parallel.py ++++ b/megatron/core/transformer/dot_product_attention_context_parallel.py +@@ -3,107 +3,12 @@ + # Some of this code was adopted from https://github.com/zhuzilin/ring-flash-attention/ + # This source code is licensed under the MIT license found in the + # LICENSE file in the root directory of this source tree. ++# Kernel is adpoted from tilelang/examples/deepseek_v32 + + import torch ++import torch.distributed as dist + from torch.nn import functional as F +- +-try: +- import einops +- +- HAVE_EINOPS = True +-except ImportError: +- HAVE_EINOPS = False +- +- +-@torch.no_grad +-def eager_attn_fwd(q, k, v, attn_bias, sinks, scale, dropout): +- """Forward pass for eager attention""" +- +- # Rearrange query, key, value to (b, h, s, d) +- b, sq, h, d = q.shape +- sk = k.shape[1] +- _q = einops.rearrange(q, 'b s h d -> b h s d') +- _k = einops.rearrange(k, 'b s h d -> b h d s') +- _v = einops.rearrange(v, 'b s h d -> b h s d') +- +- # Compute attention weights +- attn_w = torch.matmul(_q, _k) * scale +- attn_w = attn_w + attn_bias +- +- # Add sinks to attention weights +- if sinks is None: +- logits = attn_w +- else: +- _sinks = sinks.reshape(1, h, 1, 1).expand(b, -1, sq, 1) +- logits = torch.cat([attn_w, _sinks], dim=-1) +- +- # Compute attention scores +- probs = F.softmax(logits, dim=-1, dtype=logits.dtype) +- if sinks is None: +- attn_w = probs +- else: +- attn_w = probs[..., :-1] # Drop the sink +- +- # Compute attention output +- attn_output = torch.matmul(attn_w, _v) +- attn_output = einops.rearrange(attn_output, 'b h s d -> b s h d') +- attn_output = attn_output.contiguous() +- +- return attn_output, probs +- +- +-@torch.no_grad +-def eager_attn_bwd(q, k, v, attn_bias, sinks, scale, dropout, attn_output, probs, grad_output): +- """Backward pass for eager attention""" +- +- # Rearrange query, key, value to (b, h, s, d) +- b, sq, h, d = q.shape +- sk = k.shape[1] +- _q_T = einops.rearrange(q, 'b s h d -> b h d s') +- _k_T = einops.rearrange(k, 'b s h d -> b h s d') +- _v_T = einops.rearrange(v, ' b s h d -> b h d s') +- +- # Backward pass for score @ value +- if sinks is None: +- attn_w = probs +- else: +- attn_w = probs[..., :-1] # Drop the sink +- grad_output = einops.rearrange(grad_output, 'b s h d -> b h s d') +- attn_w_T = einops.rearrange(attn_w, ' b h sq sk -> b h sk sq') +- grad__v = torch.matmul(attn_w_T, grad_output) +- grad_attn_w = torch.matmul(grad_output, _v_T) +- +- # Backward pass for softmax +- if sinks is None: +- grad_probs = grad_attn_w +- else: +- dummy = torch.zeros((b, h, sq, 1), device=q.device, dtype=q.dtype) +- grad_probs = torch.cat([grad_attn_w, dummy], dim=3) +- del grad_attn_w +- grad_logits = torch._softmax_backward_data( +- grad_probs, probs, -1, probs.dtype +- ) # [b, h, sq, sk+1] +- +- # Backward pass for adding sinks +- if sinks is None: +- grad_sinks = None +- grad_attn_w = grad_logits +- else: +- grad__sinks = grad_logits[:, :, :, -1] # [b, h, sq] +- grad_sinks = einops.rearrange(grad__sinks, 'b h s -> h (b s)').sum(-1) +- grad_attn_w = grad_logits[:, :, :, :-1].contiguous() # [b, h, sq, sk] +- +- # Backward pass for q @ K^T +- grad_attn_w *= scale +- grad__q = torch.matmul(grad_attn_w, _k_T) +- grad__k = torch.matmul(_q_T, grad_attn_w) +- +- # Rearrange grads to (b, s, h, d) +- grad_v = einops.rearrange(grad__v, 'b h s d -> b s h d') +- grad_k = einops.rearrange(grad__k, 'b h d s -> b s h d') +- grad_q = einops.rearrange(grad__q, 'b h s d -> b s h d') +- return grad_q, grad_k, grad_v, grad_sinks +- ++from .tilelang_kernel import sparse_mla_bwd, sparse_mla_fwd_interface + + class AllGatherComm: + """All gather communication with async operations""" +@@ -131,212 +36,145 @@ class AllGatherComm: + handle.wait() + self.handles = [] + +- +-def to_zz_mask_attn_bias(attention_mask, cp_size, nheads, nheads_k, heads_k_stride, device, dtype): +- '''Convert the attention mask to the attention bias''' +- +- if cp_size == 1: +- zz_mask = attention_mask +- else: +- chunked = attention_mask.chunk(dim=3, chunks=cp_size * 2) +- zz_mask = [_x for _p in zip(chunked[:cp_size], reversed(chunked[cp_size:])) for _x in _p] +- zz_mask = torch.cat(zz_mask, dim=3) +- attn_bias = torch.zeros(zz_mask.shape, device=device, dtype=dtype) +- attn_bias.masked_fill_(zz_mask, float('-inf')) +- attn_bias = attn_bias.expand(-1, heads_k_stride * (nheads // nheads_k), -1, -1) +- return attn_bias +- +- + class AttentionFuncionWithContextParallel(torch.autograd.Function): + """Native attention function with context parallelism.""" + ++ # q: [seq_len_shard, batch, nheads, dim] ++ # k: [seq_len_kv_shard, batch, 1, dim] ++ # v: [seq_len_kv_shard, batch, 1, dim_v] ++ # indices: [batch, 1, seq_len, topk] ++ # masks: [batch, 1, seq_len, seq_len_kv] + @staticmethod +- def forward(ctx, q, k, v, attention_mask, attention_dropout, softmax_scale, pg): ++ def forward(ctx, q, k, dim_v, indices, masks, attention_dropout, softmax_scale, pg): + '''Forward pass for the native attention function with context parallelism''' + +- # Assert einops exists +- if not HAVE_EINOPS: +- raise ImportError("einops is required by the attention CP but cannot be imported.") +- +- # Initialize communication group and constants + cp_size = 1 + if pg is not None: + cp_size = torch.distributed.get_world_size(pg) + comm = AllGatherComm(group=pg) +- nheads = q.shape[2] +- nheads_k = k.shape[2] +- heads_k_stride = 1 +- assert nheads % nheads_k == 0 and nheads_k % heads_k_stride == 0 +- outs = [] +- probs = [] +- +- # Initialize KV buffers +- kv_buffer = torch.empty( +- (2, k.shape[0] * cp_size, k.shape[1], heads_k_stride, k.shape[3]), ++ ++ k_buffer = torch.empty( ++ (k.shape[0] * cp_size, k.shape[1], 1, k.shape[3]), + dtype=k.dtype, + device=k.device, + ) +- kv_buffer_copy = torch.empty_like(kv_buffer) +- +- # All-gather first chunk of KV buffers +- k_0 = k[:, :, :heads_k_stride].contiguous() +- v_0 = v[:, :, :heads_k_stride].contiguous() +- comm.all_gather(kv_buffer_copy[0], k_0) +- comm.all_gather(kv_buffer_copy[1], v_0) +- +- # Prepare attention bias +- attn_bias = to_zz_mask_attn_bias( +- attention_mask, cp_size, nheads, nheads_k, heads_k_stride, q.device, q.dtype +- ) +- +- # Iterate over heads +- for i in range(0, nheads_k, heads_k_stride): +- # Wait for previous all-gather to complete +- comm.wait() +- kv_buffer, kv_buffer_copy = kv_buffer_copy, kv_buffer +- # All-gather the next portion of KV buffers if not the last iteration +- if i < nheads_k - heads_k_stride: +- kvsl = i + heads_k_stride +- kvsr = kvsl + heads_k_stride +- send_k = k[:, :, kvsl:kvsr].contiguous() +- send_v = v[:, :, kvsl:kvsr].contiguous() +- comm.all_gather(kv_buffer_copy[0], send_k) +- comm.all_gather(kv_buffer_copy[1], send_v) +- +- # Prepare query, key, value for attention +- q_i = q[:, :, i * nheads // nheads_k : (i + heads_k_stride) * nheads // nheads_k] +- k_i = kv_buffer[0] +- v_i = kv_buffer[1] +- +- # Rearrange query, key, value to (b, s, h, d) +- q_i = einops.rearrange(q_i, 's b h d -> b s h d') +- k_i = einops.rearrange(k_i, 's b h d -> b s h d') +- v_i = einops.rearrange(v_i, 's b h d -> b s h d') +- +- # Forward pass +- out_i, probs_i = eager_attn_fwd( +- q_i, k_i, v_i, attn_bias, None, softmax_scale, attention_dropout +- ) +- outs.append(out_i) +- probs.append(probs_i) +- +- # Concatenate outputs and rearrange to (s, b, h, d) +- out = torch.cat(outs, dim=2) +- out = einops.rearrange(out, 'b s h d -> s b h d') +- +- # Save contexts for backward pass +- ctx.save_for_backward(q, k, v, attention_mask, *outs, *probs) ++ comm.all_gather(k_buffer, k) ++ comm.wait() ++ ++ zz_indices = indices.transpose(1, 2) ++ zz_masks = masks.transpose(1, 2) ++ ++ q_i = q ++ k_i = k_buffer ++ ++ s_, b_, h_, d_ = q_i.shape ++ q_i = q_i.transpose(0, 1).flatten().view(b_, s_, h_, d_) ++ s_, b_, h_, d_ = k_i.shape ++ k_i = k_i.transpose(0, 1).flatten().view(b_, s_, h_, d_) ++ zz_indices_i = zz_indices ++ b_, s_, g_, topk_ = zz_indices_i.shape ++ zz_indices_i = zz_indices_i.flatten().view(b_, s_, g_, topk_) ++ zz_masks_i = zz_masks ++ b_, s_, g_, skv_ = zz_masks_i.shape ++ zz_masks_i = zz_masks_i.flatten().view(b_, s_, g_, skv_) ++ ++ out_i, lse_i = sparse_mla_fwd_interface(q_i.contiguous(), k_i, zz_indices_i, zz_masks_i, dim_v, sm_scale = softmax_scale) ++ ++ # out: [B, seq_len_shard, h, dim] -> [seq_len, B, h, dim] ++ b_, s_, h_, d_ = out_i.shape ++ out_i = out_i.transpose(0, 1).flatten().view(s_, b_, h_, d_).contiguous() ++ ++ # outs: [[B, seq_len_shard, nheads // kv_group, dim], ...., [B, seq_len_shard, nheads // kv_group, dim]], repeat kv_group // heads_kv_stride times ++ # lses: [[B, seq_len_shard, heads_kv_stride], ...., [B, seq_len_shard, heads_kv_stride]], repeat kv_group // heads_kv_stride times ++ ctx.save_for_backward(q, k, indices, masks, out_i, lse_i) + ctx.dropout = attention_dropout +- ctx.scale = softmax_scale +- ctx.heads_k_stride = heads_k_stride # TODO make it configurable ++ ctx.softmax_scale = softmax_scale ++ ctx.dim_v = dim_v + ctx.pg = pg + +- return out ++ return out_i + + @staticmethod + def backward(ctx, dout): + '''Backward pass for the native attention function with context parallelism''' + +- # Initialize or resume constants and communication group +- q, k, v, attention_mask, *rest = ctx.saved_tensors +- nheads = q.shape[2] +- nheads_k = k.shape[2] +- heads_k_stride = ctx.heads_k_stride +- assert nheads_k % heads_k_stride == 0 +- outs = rest[: nheads_k // heads_k_stride] +- probs = rest[nheads_k // heads_k_stride :] ++ q, k, indices, masks, out, lse = ctx.saved_tensors ++ s, b, heads, dim = q.shape ++ dim_v = ctx.dim_v ++ softmax_scale = ctx.softmax_scale ++ + pg = ctx.pg + cp_size = 1 + if pg is not None: + cp_size = torch.distributed.get_world_size(pg) + comm = AllGatherComm(group=pg) + +- # Initialize KV buffers +- kv_buffer = torch.empty( +- (2, k.shape[0] * cp_size, k.shape[1], heads_k_stride, k.shape[3]), ++ k_buffer = torch.empty( ++ (k.shape[0] * cp_size, k.shape[1], 1, k.shape[3]), + dtype=k.dtype, + device=k.device, + ) +- kv_buffer_copy = torch.empty_like(kv_buffer) +- +- # All-gather first chunk of KV buffers +- dq = [] +- dk = [] +- dv = [] +- k_0 = k[:, :, :heads_k_stride].contiguous() +- v_0 = v[:, :, :heads_k_stride].contiguous() +- comm.all_gather(kv_buffer_copy[0], k_0) +- comm.all_gather(kv_buffer_copy[1], v_0) +- +- # Prepare attention bias +- attn_bias = to_zz_mask_attn_bias( +- attention_mask, cp_size, nheads, nheads_k, heads_k_stride, q.device, q.dtype +- ) + +- # Iterate over heads +- for i in range(0, nheads_k, heads_k_stride): +- # Slice query and output for this iteration +- q_slice = slice(i * nheads // nheads_k, (i + heads_k_stride) * nheads // nheads_k) +- q_i = q[:, :, q_slice] +- dout_i = dout[:, :, q_slice] +- +- # Wait for previous all-gather to complete +- comm.wait() +- kv_buffer, kv_buffer_copy = kv_buffer_copy, kv_buffer +- +- # All-gather the next portion of KV buffers if not the last iteration +- if i < nheads_k - heads_k_stride: +- kvsl = i + heads_k_stride +- kvsr = kvsl + heads_k_stride +- send_k = k[:, :, kvsl:kvsr].contiguous() +- send_v = v[:, :, kvsl:kvsr].contiguous() +- comm.all_gather(kv_buffer_copy[0], send_k) +- comm.all_gather(kv_buffer_copy[1], send_v) +- +- # Prepare key, value for attention +- k_i = kv_buffer[0] +- v_i = kv_buffer[1] +- +- # Rearrange query, key, value to (b, s, h, d) +- q_i = einops.rearrange(q_i, 's b h d -> b s h d') +- k_i = einops.rearrange(k_i, 's b h d -> b s h d') +- v_i = einops.rearrange(v_i, 's b h d -> b s h d') +- dout_i = einops.rearrange(dout_i, 's b h d -> b s h d') +- +- # Backward pass +- dq_i, _dk_i, _dv_i, _ = eager_attn_bwd( +- q_i, k_i, v_i, attn_bias, None, ctx.scale, ctx.dropout, outs[i], probs[i], dout_i +- ) ++ comm.all_gather(k_buffer, k) ++ comm.wait() ++ ++ zz_indices = indices.transpose(1, 2) ++ zz_masks = masks.transpose(1, 2) ++ ++ k_i = k_buffer ++ ++ dq_list = [] ++ dk_list = [] ++ ++ s_, b_, h_, d_ = q.shape ++ q = q.transpose(0, 1).flatten().view(b_, s_, h_, d_) ++ s_, b_, h_, d_ = k_i.shape ++ k_i = k_i.transpose(0, 1).flatten().view(b_, s_, h_, d_) ++ s_, b_, h_, d_ = dout.shape ++ dout = dout.transpose(0, 1).flatten().view(b_, s_, h_, d_) ++ s_, b_, h_, d_ = out.shape ++ out = out.transpose(0, 1).flatten().view(b_, s_, h_, d_) ++ b_, s_, h_ = lse.shape ++ lse = lse.flatten().view(b_, s_, h_) ++ zz_indices_i = zz_indices ++ b_, s_, g_, topk_ = zz_indices_i.shape ++ zz_indices_i = zz_indices_i.flatten().view(b_, s_, g_, topk_) ++ zz_masks_i = zz_masks ++ b_, s_, g_, skv_ = zz_masks_i.shape ++ zz_masks_i = zz_masks_i.flatten().view(b_, s_, g_, skv_) ++ ++ heads_kv_stride = 16 ++ for i in range(0, heads, heads_kv_stride): ++ q_slice = slice(i, min(i + heads_kv_stride, heads)) ++ q_i = q[:, :, q_slice, :].contiguous() ++ dout_i = dout[:, :, q_slice, :].contiguous() ++ out_i = out[:, :, q_slice, :].contiguous() ++ lse_i = lse[:, :, q_slice].contiguous() ++ ++ # TODO: needs casual = True, may not be compatible with zz ++ dq_i, _dk_i = sparse_mla_bwd(q_i, k_i, out_i, dout_i, zz_indices_i, zz_masks_i, lse_i, dim_v, sm_scale = softmax_scale) ++ ++ b_, s_, h_, d_ = dq_i.shape ++ dq_i = dq_i.transpose(0, 1).flatten().view(s_, b_, h_, d_).contiguous() ++ b_, s_, h_, d_ = _dk_i.shape ++ _dk_i = _dk_i.transpose(0, 1).flatten().view(s_, b_, h_, d_).contiguous() + +- # Rearrange gradients to (s, b, h, d) +- dq_i = einops.rearrange(dq_i, 'b s h d -> s b h d') +- _dk_i = einops.rearrange(_dk_i, 'b s h d -> s b h d') +- _dv_i = einops.rearrange(_dv_i, 'b s h d -> s b h d') + if pg is None: + dk_i = _dk_i +- dv_i = _dv_i + else: +- # Reduce-scatter gradients if CP > 1 + dk_i = torch.zeros( + (k_i.shape[1] // cp_size, k_i.shape[0], k_i.shape[2], k_i.shape[3]), + device=k_i.device, +- dtype=k_i.dtype, +- ) +- dv_i = torch.zeros( +- (v_i.shape[1] // cp_size, v_i.shape[0], v_i.shape[2], v_i.shape[3]), +- device=v_i.device, +- dtype=v_i.dtype, ++ dtype=torch.float32, + ) + torch.distributed.reduce_scatter_tensor(dk_i, _dk_i, group=pg) +- torch.distributed.reduce_scatter_tensor(dv_i, _dv_i, group=pg) + +- # Collect gradients +- dq.append(dq_i) +- dk.append(dk_i) +- dv.append(dv_i) ++ dq_list.append(dq_i) ++ dk_list.append(dk_i) + + # Concatenate gradients and return +- dq = torch.cat(dq, dim=2) +- dk = torch.cat(dk, dim=2) +- dv = torch.cat(dv, dim=2) +- return dq, dk, dv, None, None, None, None ++ dq = torch.cat(dq_list, dim=2) ++ dk_ = torch.cat(dk_list, dim=2) ++ dk = torch.sum(dk_, dim=2, keepdim=True).to(torch.bfloat16) ++ ++ return dq, dk, None, None, None, None, None, None +\ No newline at end of file +diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py +index 353b31e9b..221e93500 100644 +--- a/megatron/core/transformer/experimental_attention_variant/dsa.py ++++ b/megatron/core/transformer/experimental_attention_variant/dsa.py +@@ -6,6 +6,7 @@ from dataclasses import dataclass + from typing import Optional, Tuple, Union + + import torch ++import einops + + from megatron.core import parallel_state + from megatron.core.models.common.embeddings import ( +@@ -20,6 +21,7 @@ from megatron.core.transformer.enums import AttnMaskType + from megatron.core.transformer.module import MegatronModule + from megatron.core.transformer.spec_utils import ModuleSpec, build_module + from megatron.core.transformer.transformer_config import TransformerConfig ++from megatron.core.transformer.dot_product_attention_context_parallel import AllGatherComm, AttentionFuncionWithContextParallel + + try: + from fast_hadamard_transform import hadamard_transform +@@ -191,44 +193,72 @@ def compute_dsa_indexer_loss( + Returns: + index_loss: KL divergence loss (scalar). + """ +- sq, b, np, hn = query.size() +- sk = key.size(0) ++ cp_size = parallel_state.get_context_parallel_world_size() + +- # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn] +- query = query.permute(1, 2, 0, 3).reshape(b * np, sq, hn) +- # [sk, b, np, hn] -> [b, np, hn, sk] -> [b * np, hn, sk] +- key = key.permute(1, 2, 3, 0).reshape(b * np, hn, sk) +- # Compute attention scores [b * np, sq, sk] +- attention_scores = torch.bmm(query.float(), key.float()) * softmax_scale +- # Reshape to [b, np, sq, sk] +- attention_scores = attention_scores.reshape(b, np, sq, sk) ++ if cp_size > 1: ++ sq_local, b, np, hn = query.size() ++ sk_local = key.size(0) ++ sk_global = sk_local * cp_size + +- # causal_mask [sq, sk] +- causal_mask = torch.triu( +- torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device), +- diagonal=1, +- ) +- # index_mask [b, sq, sk] +- index_mask = torch.full( +- (b, sq, sk), float("-inf"), dtype=torch.float32, device=causal_mask.device +- ).scatter_(-1, topk_indices, 0) +- +- # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv] +- attention_scores += causal_mask.view(1, 1, sq, sk) +- if sparse_loss: +- # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk] +- attention_scores += index_mask.view(b, 1, sq, sk) +- # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk] +- index_scores += index_mask +- +- # [b, np, sq, sk] -> [b, np, sq, sk] +- attention_scores = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32) +- # [b, sq, sk] -> [b, sq, sk] +- index_scores = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) ++ causal_mask = get_causal_mask(sq_local, sk_local, query.device) ++ float_mask = torch.zeros_like(causal_mask, dtype=torch.float32).masked_fill( ++ causal_mask, float('-inf') ++ ) ++ ++ index_mask = torch.full( ++ (b, sq_local, sk_global), float("-inf"), dtype=torch.float32, device=causal_mask.device ++ ).scatter_(-1, topk_indices, 0) + +- # Sum attention scores across heads. +- # [batch, heads, seqlen_q, seqlen_k] -> [batch, seqlen_q, seqlen_k] +- attention_scores = attention_scores.sum(dim=1) ++ float_mask = float_mask.view(1, 1, sq_local, sk_global) ++ float_mask = index_mask.view(b, 1, sq_local, sk_global) + float_mask if sparse_loss else float_mask ++ ++ # because the attention computation is more heavy in memory (has head dim), ++ # we apply cp (all-gather backend) on attention scores computation ++ attention_scores = compute_attention_scores_with_cp(query, key, float_mask, softmax_scale) # [b, sq_local, sk_global] ++ ++ index_scores = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) ++ ++ else: ++ sq, b, np, hn = query.size() ++ sk = key.size(0) ++ ++ # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn] ++ query = query.permute(1, 2, 0, 3).reshape(b * np, sq, hn) ++ # [sk, b, np, hn] -> [b, np, hn, sk] -> [b * np, hn, sk] ++ key = key.permute(1, 2, 3, 0).reshape(b * np, hn, sk) ++ # Compute attention scores [b * np, sq, sk] ++ attention_scores = torch.bmm(query.float(), key.float()) * softmax_scale ++ # Reshape to [b, np, sq, sk] ++ attention_scores = attention_scores.reshape(b, np, sq, sk) ++ ++ # causal_mask [sq, sk] ++ causal_mask = torch.triu( ++ torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device), ++ diagonal=1, ++ ) ++ # index_mask [b, sq, sk] ++ index_mask = torch.full( ++ (b, sq, sk), float("-inf"), dtype=torch.float32, device=causal_mask.device ++ ).scatter_(-1, topk_indices, 0) ++ ++ # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv] ++ attention_scores += causal_mask.view(1, 1, sq, sk) ++ if sparse_loss: ++ # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk] ++ attention_scores += index_mask.view(b, 1, sq, sk) ++ # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk] ++ index_scores += index_mask ++ ++ # [b, np, sq, sk] -> [b, np, sq, sk] ++ attention_scores = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32) ++ # [b, sq, sk] -> [b, sq, sk] ++ index_scores = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) ++ ++ # Sum attention scores across heads. ++ # [batch, heads, seqlen_q, seqlen_k] -> [batch, seqlen_q, seqlen_k] ++ attention_scores = attention_scores.sum(dim=1) ++ ++ # Common part + if pg_collection.tp.size() > 1: + # attention scores are scattered to TP ranks in head dimension. + torch.distributed.all_reduce(attention_scores.contiguous(), group=pg_collection.tp) +@@ -251,6 +281,56 @@ def compute_dsa_indexer_loss( + + return indexer_loss + ++def compute_attention_scores_with_cp(q, k, attn_bias, scale, heads_k_stride = 1): ++ """ ++ compute attention scores of q_local @ k_global with CP all-gather backend ++ parallel on n_heads dimension ++ """ ++ pg = parallel_state.get_context_parallel_group() ++ cp_size = parallel_state.get_context_parallel_world_size() ++ ++ sq_local, b, nheads, hn_q = q.shape ++ sk_local, _, nheads_k, hn_k = k.shape ++ sk_global = sk_local * cp_size ++ ++ assert nheads % nheads_k == 0 and nheads_k % heads_k_stride == 0 ++ ++ comm = AllGatherComm(group=pg) ++ attns = torch.zeros(b, heads_k_stride, sq_local, sk_global, dtype=q.dtype, device=q.device) ++ ++ k_buffer = torch.empty( ++ (sk_global, b, heads_k_stride, hn_k), ++ dtype=k.dtype, ++ device=k.device ++ ) ++ k_buffer_copy = torch.empty_like(k_buffer) ++ k_0 = k[:, :, :heads_k_stride].contiguous() ++ comm.all_gather(k_buffer_copy, k_0) ++ ++ attn_bias = attn_bias.expand(-1, heads_k_stride * (nheads // nheads_k), -1, -1) ++ ++ for i in range(0, nheads_k, heads_k_stride): ++ comm.wait() ++ k_buffer, k_buffer_copy = k_buffer_copy, k_buffer ++ if i < nheads_k - heads_k_stride: ++ kvsl = i + heads_k_stride ++ kvsr = kvsl + heads_k_stride ++ send_k = k[:, :, kvsl:kvsr].contiguous() ++ comm.all_gather(k_buffer_copy, send_k) ++ q_i = q[:, :, i * nheads // nheads_k : (i + heads_k_stride) * nheads // nheads_k] ++ k_i = k_buffer ++ ++ _q_i = einops.rearrange(q_i, 's b h d -> b h s d') ++ _k_i = einops.rearrange(k_i, 's b h d -> b h d s') ++ attn_i = torch.matmul(_q_i.float(), _k_i.float()) * scale + attn_bias ++ attn_i = torch.nn.functional.softmax(attn_i, dim=-1, dtype=torch.float32) ++ ++ attns = attns + attn_i ++ ++ attns = torch.sum(attns, dim=1) ++ ++ return attns ++ + + class DSAIndexerLossAutoScaler(torch.autograd.Function): + """An AutoScaler that triggers the backward pass and scales the grad for indexer loss. +@@ -496,7 +576,15 @@ class DSAIndexer(MegatronModule): + # Compute attention scores: q @ k^T + # [seqlen_q, batch, index_n_heads, index_head_dim] @ [seqlen_k, batch, index_head_dim]^T + # -> [seqlen_q, batch, index_n_heads, seqlen_k] +- index_scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float()) ++ cp_size = parallel_state.get_context_parallel_world_size() ++ if cp_size == 1: ++ index_scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float()) ++ else: ++ # because k is small (only 1 head), do just one all_gather ++ k_buffer = torch.cat(torch.distributed.nn.functional.all_gather(k, group=self.pg_collection.cp), dim=0) # k_buffer: [[chunk_0, chunk_3, chunk_1, chunk_2], batch, index_head_dim] ++ index_scores = torch.einsum('sbhd,tbd->sbht', q.float(), k_buffer.float()) # [s_q_local, batch, index_n_heads, s_k_global] ++ # rank 0: q [chunk_0, chunk_3], k[chunk_0, chunk_3, chunk_1, chunk_2] ++ # rank 1: q [chunk_1, chunk_2], k[chunk_0, chunk_3, chunk_1, chunk_2] + + # Apply ReLU activation. + index_scores = torch.relu(index_scores) +@@ -546,14 +634,10 @@ class DSAIndexer(MegatronModule): + None, None, x, self.config, packed_seq_params + ) + if self.config.rope_type == "rope": +- rotary_pos_emb = self.rotary_pos_emb( +- rotary_seq_len, packed_seq_params=packed_seq_params +- ) ++ rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq_params=packed_seq_params) + mscale = 1.0 + else: +- rotary_pos_emb, mscale = self.rotary_pos_emb( +- rotary_seq_len, packed_seq_params=packed_seq_params +- ) ++ rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq_params=packed_seq_params) + + # ========================================= + # Gather inputs if sp is enabled +@@ -610,7 +694,9 @@ class DSAIndexer(MegatronModule): + # ========================================= + # Select top-k indices + # ========================================= +- topk_k = min(self.index_topk, seqlen) ++ cp_size = parallel_state.get_context_parallel_world_size() ++ seqlen_k_global = k.shape[0] * cp_size ++ topk_k = min(self.index_topk, seqlen_k_global) + # [batch, seqlen, index_topk] + topk_indices = index_scores.topk(topk_k, dim=-1)[1] + +@@ -691,6 +777,48 @@ def unfused_dsa_fn(query, key, value, topk_indices, softmax_scale): + output = output.reshape(sq, b, np * hnv) + return output + ++def get_causal_mask(sq, skv, device): ++ cp_size = parallel_state.get_context_parallel_world_size() ++ cp_rank = parallel_state.get_context_parallel_rank() ++ skv_global = skv * cp_size ++ ++ if cp_size == 1: ++ causal_mask = torch.triu( ++ torch.ones((sq, skv), dtype=torch.bool, device=device), ++ diagonal=1, ++ ) ++ else: ++ sq_half = sq // 2 ++ global_q_positions = torch.cat([ ++ torch.arange(cp_rank * sq_half, (cp_rank + 1) * sq_half, device=device), ++ torch.arange(skv_global - (cp_rank + 1) * sq_half, skv_global - cp_rank * sq_half, device=device) ++ ]) ++ ++ global_k_positions = torch.arange(skv_global, device=device) ++ # [sq, 1] < [1, skv_global] -> [sq, skv_global] ++ causal_mask = global_q_positions.unsqueeze(1) < global_k_positions.unsqueeze(0) ++ # convert to zz mask ++ chunked = causal_mask.chunk(dim=1, chunks=cp_size * 2) ++ causal_mask = [_x for _p in zip(chunked[:cp_size], reversed(chunked[cp_size:])) for _x in _p] ++ causal_mask = torch.cat(causal_mask, dim=1) ++ ++ return causal_mask ++ ++def unfused_dsa_fn_with_cp(query, key, dim_v, topk_indices, softmax_scale): ++ pg = parallel_state.get_context_parallel_group() ++ sq, b, np, hn = query.size() ++ skv = key.size(0) ++ ++ topk = topk_indices.shape[-1] ++ topk_indices = topk_indices.unsqueeze(1) ++ topk_indices = topk_indices.expand(-1, key.shape[2], -1, -1).contiguous().to(torch.int32) ++ causal_masks = get_causal_mask(sq, skv, query.device) ++ causal_masks = causal_masks[None, None, :, :] ++ causal_masks = causal_masks.expand(b, key.shape[2], -1, -1).contiguous() ++ output = AttentionFuncionWithContextParallel.apply( ++ query, key, dim_v, topk_indices, causal_masks, 0.0, softmax_scale, pg ++ ) ++ return output.reshape(sq, b, np * dim_v) + + class DSAttention(MegatronModule): + """ +@@ -733,7 +861,6 @@ class DSAttention(MegatronModule): + self, + query: torch.Tensor, + key: torch.Tensor, +- value: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + attention_mask: torch.Tensor, +@@ -747,7 +874,6 @@ class DSAttention(MegatronModule): + Args: + query: Query tensor [sq, b, np, hn]. + key: Key tensor [skv, b, np, hn]. +- value: Value tensor [skv, b, np, hnv]. + x: Original hidden states [sq, b, hidden_size]. + qr: Low-rank query representation [sq, b, q_lora_rank]. + attention_mask: Attention mask tensor [b, 1, sq, sk]. +@@ -758,9 +884,11 @@ class DSAttention(MegatronModule): + Returns: + output: Output tensor [sq, b, hidden_size] + """ +- sq, b, np, hn = query.size() +- skv = key.size(0) +- hnv = value.size(3) ++ dim_v = self.config.kv_lora_rank ++ # torch.Size([128, 1, 64, 576]) ++ sq, b, nheads, dim = query.size() ++ # torch.Size([128, 1, 1, 576]) ++ skv, _, kv_groups, _ = key.shape + + # Detach x and qr to prevent gradients of indexer from flowing back to the main model. + x = x.detach() +@@ -772,18 +900,17 @@ class DSAttention(MegatronModule): + # Generate upper triangular mask with -inf above diagonal, 0 elsewhere + # torch.triu with diagonal=1 creates upper triangular matrix (excluding main diagonal) + # float_mask [sq, skv] +- float_mask = torch.triu( +- torch.full((sq, skv), float('-inf'), dtype=torch.float32, device=x.device), +- diagonal=1, +- ) ++ mask = get_causal_mask(sq, skv, x.device) + else: +- assert attention_mask.shape == (b, 1, sq, skv), 'attention_mask shape mismatch' ++ skv_global = skv * parallel_state.get_context_parallel_world_size() ++ assert attention_mask.shape == (b, 1, sq, skv_global), 'attention_mask shape mismatch' + # [b, 1, sq, skv] -> [b, sq, skv] + mask = attention_mask.squeeze() +- # float_mask [b, sq, skv] +- float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill( +- mask, float('-inf') +- ) ++ ++ # float_mask [b, sq, skv] ++ float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill( ++ mask, float('-inf') ++ ) + + # =================================== + # Get index scores and top-k indices +@@ -795,32 +922,6 @@ class DSAttention(MegatronModule): + # =================================== + # Run sparse attention kernel + # =================================== +- output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale) +- +- # =================================== +- # Attach indexer loss +- # =================================== +- if self.training and torch.is_grad_enabled(): +- # Compute KL divergence loss between indexer scores and true attention scores +- indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) +- indexer_loss = compute_dsa_indexer_loss( +- index_scores, +- topk_indices, +- query.detach(), +- key.detach(), +- self.softmax_scale, +- indexer_loss_coeff, +- getattr(self.config, "dsa_indexer_use_sparse_loss", False), +- self.indexer.pg_collection, +- ) +- # Save indexer loss for logging +- if indexer_loss_coeff > 0: +- DSAIndexerLossLoggingHelper.save_loss_to_tracker( +- loss=indexer_loss, +- layer_number=self.layer_number, +- num_layers=self.config.num_layers, +- ) +- # Attach loss to output +- output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) +- ++ output = unfused_dsa_fn_with_cp(query, key, dim_v, topk_indices, self.softmax_scale) ++ + return output +diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py +index 28cff06f5..befb5c124 100644 +--- a/megatron/core/transformer/moe/moe_utils.py ++++ b/megatron/core/transformer/moe/moe_utils.py +@@ -586,6 +586,9 @@ def topk_routing_with_score_function( + ) + else: + return torch.topk(scores, k=topk, dim=1) ++ ++ from miles.utils.routing_replay import get_routing_replay_compute_topk ++ compute_topk = get_routing_replay_compute_topk(compute_topk) + + if score_function == "softmax": + if use_pre_softmax: +diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py +index 16fc9d9af..3c50a9516 100644 +--- a/megatron/core/transformer/moe/router.py ++++ b/megatron/core/transformer/moe/router.py +@@ -200,6 +200,9 @@ class TopKRouter(Router): + else: + self.global_tokens_per_expert = None + self.ga_steps = None ++ ++ from miles.utils.routing_replay import register_routing_replay ++ register_routing_replay(self) + + def _maintain_float32_expert_bias(self): + """ +diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py +index ed90fdffa..7a7597d66 100644 +--- a/megatron/core/transformer/multi_latent_attention.py ++++ b/megatron/core/transformer/multi_latent_attention.py +@@ -15,7 +15,7 @@ except ImportError: + HAVE_EINOPS = False + + +-from megatron.core import tensor_parallel ++from megatron.core import parallel_state, tensor_parallel + from megatron.core.models.common.embeddings import ( + RotaryEmbedding, + YarnRotaryEmbedding, +@@ -312,7 +312,6 @@ class MultiLatentAttention(Attention): + core_attn_out = self.core_attention( + query, + key, +- value, + x=hidden_states, + qr=q_compressed, + attention_mask=attention_mask, +@@ -371,6 +370,19 @@ class MultiLatentAttention(Attention): + self.qkv_up_checkpoint.discard_output_and_register_recompute(core_attn_out) + self.qkv_up_checkpoint = None + ++ s_, b_ = core_attn_out.size(0), core_attn_out.size(1) ++ core_attn_out = core_attn_out.view( ++ s_, b_, ++ self.num_attention_heads_per_partition, ++ self.config.kv_lora_rank ++ ) ++ ++ # einsum: "sbhk,hdk->sbhd" ++ core_attn_out = torch.einsum("sbhk,hdk->sbhd", core_attn_out, self.up_v_weight_) ++ core_attn_out = core_attn_out.contiguous() ++ core_attn_out = core_attn_out.view(s_, b_, -1) ++ core_attn_out = core_attn_out.contiguous() ++ + # ================= + # Output. [sq, b, h] + # ================= +@@ -555,11 +567,7 @@ class MLASelfAttention(MultiLatentAttention): + assert ( + hidden_states.ndim == 3 + ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D" +- if packed_seq_params is not None: +- assert ( +- packed_seq_params.local_cp_size is None +- ), "hybrid_context_parallel is not supported with MLA yet and is planned for future. \ +- Please disable hybrid_context_parallel." ++ + + inference_context = deprecate_inference_params(inference_context, inference_params) + +@@ -576,9 +584,7 @@ class MLASelfAttention(MultiLatentAttention): + rotary_pos_sin = None + packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + if self.config.rope_type == "rope": +- rotary_pos_emb = self.rotary_pos_emb( +- rotary_seq_len, packed_seq_params=packed_seq_params +- ) ++ rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq_params=packed_seq_params) + else: + if self.config.apply_rope_fusion: + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( +@@ -591,11 +597,9 @@ class MLASelfAttention(MultiLatentAttention): + and fused_apply_mla_rope_for_kv is not None + ), "Fused MLA RoPE apply is not imported successfully" + else: +- rotary_pos_emb, mscale = self.rotary_pos_emb( +- rotary_seq_len, packed_seq_params=packed_seq_params +- ) ++ rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq_params=packed_seq_params) + +- if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': ++ if packed_seq_params is not None: + if packed_seq_params.cu_seqlens_q_padded is not None: + cu_seqlens_q = packed_seq_params.cu_seqlens_q_padded + else: +@@ -867,6 +871,98 @@ class MLASelfAttention(MultiLatentAttention): + + return query, key, value + ++ def mla_absorb(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb): ++ if self.config.q_lora_rank is not None: ++ # q_compressed: [num_tokens, q_lora_rank] ++ # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)] ++ q, _ = self.linear_q_up_proj(q_compressed) ++ else: ++ # q_compressed: [num_tokens, hidden_size] ++ # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)] ++ q, _ = self.linear_q_proj(q_compressed) ++ ++ # q: [num_tokens, n, q_head_dim] ++ q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim) ++ ++ # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] ++ k_pos_emb = torch.unsqueeze(k_pos_emb, -2) ++ ++ if self.config.apply_rope_fusion: ++ raise NotImplementedError( ++ "RoPE fusion is not yet supported with absorption training. " ++ "Please set apply_rope_fusion=False." ++ ) ++ else: ++ q_len = q.size()[0] ++ if inference_context is not None: ++ # add offset to the sequence start for inference ++ sequence_start = inference_context.sequence_len_offset ++ sequence_end = sequence_start + q_len ++ rotary_pos_emb = rotary_pos_emb[sequence_start:sequence_end] ++ elif packed_seq_params is None or self.config.context_parallel_size == 1: ++ rotary_pos_emb = rotary_pos_emb[0:q_len] ++ ++ # q_no_pe: [num_tokens, n, qk_head_dim] ++ # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] ++ q_no_pe, q_pos_emb = torch.split( ++ q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1 ++ ) ++ ++ # q_no_pe: [num_tokens, n, qk_head_dim] ++ # up_k_weight: [n, qk_head_dim, kv_lora_rank] ++ # q_absorbed: [num_tokens, n, kv_lora_rank] ++ q_absorbed = torch.einsum("...hd,hdk->...hk", q_no_pe, self.up_k_weight_) ++ ++ # TODO: Does it match ZZ? SP does not need but CP needs ++ if self.config.sequence_parallel: ++ kv_compressed = gather_from_sequence_parallel_region(kv_compressed, group=self.tp_group) ++ ++ # kv_compressed: [num_tokens, kv_lora_rank] ++ if kv_compressed.ndim == 3: # [s, b, kv_lora_rank] ++ k_content = kv_compressed.unsqueeze(2).expand( ++ -1, -1, 1, -1 ++ ) ++ else: # [t, kv_lora_rank] for packed sequence ++ k_content = kv_compressed.unsqueeze(1).expand( ++ -1, 1, -1 ++ ) ++ ++ # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] ++ q_pos_emb = apply_rotary_pos_emb( ++ q_pos_emb, ++ rotary_pos_emb, ++ config=self.config, ++ cu_seqlens=cu_seqlens_q, ++ mscale=mscale, ++ cp_group=self.pg_collection.cp, ++ ) ++ # k_pos_emb: [num_tokens, 1, qk_pos_emb_head_dim] ++ k_pos_emb = apply_rotary_pos_emb( ++ k_pos_emb, ++ rotary_pos_emb, ++ config=self.config, ++ cu_seqlens=cu_seqlens_kv, ++ mscale=mscale, ++ cp_group=self.pg_collection.cp, ++ ) ++ ++ # query: [num_tokens, n, kv_lora_rank + qk_pos_emb_head_dim] ++ query = torch.cat([q_absorbed, q_pos_emb], dim=-1) ++ ++ # key: [num_tokens, n, kv_lora_rank + qk_pos_emb_head_dim] ++ if k_pos_emb.ndim == 4: ++ k_pos_emb = k_pos_emb.expand(-1, -1, 1, -1) ++ else: ++ assert k_pos_emb.ndim == 3 ++ k_pos_emb = k_pos_emb.expand(-1, 1, -1) ++ ++ key = torch.cat([k_content, k_pos_emb], dim=-1) ++ ++ query = query.contiguous() ++ key = key.contiguous() ++ ++ return query, key ++ + if self.recompute_up_proj: + quantization = self.config.fp8 or self.config.fp4 + self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput(fp8=quantization) +@@ -882,9 +978,10 @@ class MLASelfAttention(MultiLatentAttention): + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) + else: +- query, key, value = qkv_up_proj_and_rope_apply( ++ query, key = mla_absorb( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) ++ value = None + + if return_compressed_tensors: + return query, key, value, q_compressed, kv_compressed +@@ -1128,3 +1225,26 @@ class MLASelfAttention(MultiLatentAttention): + ) + + return weight_kv_updated ++ ++ @property ++ def up_k_weight_(self): ++ # linear_kv_up_proj.weight: [num_heads_per_partition * (qk_head_dim + v_head_dim), kv_lora_rank] ++ weight = self.linear_kv_up_proj.weight ++ weight_reshaped = weight.view( ++ self.num_attention_heads_per_partition, ++ self.config.qk_head_dim + self.config.v_head_dim, ++ self.config.kv_lora_rank, ++ ) ++ # [num_heads_per_partition, qk_head_dim, kv_lora_rank] ++ return weight_reshaped[:, :self.config.qk_head_dim, :] ++ ++ @property ++ def up_v_weight_(self): ++ weight = self.linear_kv_up_proj.weight ++ weight_reshaped = weight.view( ++ self.num_attention_heads_per_partition, ++ self.config.qk_head_dim + self.config.v_head_dim, ++ self.config.kv_lora_rank, ++ ) ++ # [num_heads_per_partition, v_head_dim, kv_lora_rank] ++ return weight_reshaped[:, self.config.qk_head_dim:, :] +\ No newline at end of file +diff --git a/megatron/core/transformer/tilelang_kernel/__init__.py b/megatron/core/transformer/tilelang_kernel/__init__.py +new file mode 100644 +index 000000000..c63794256 +--- /dev/null ++++ b/megatron/core/transformer/tilelang_kernel/__init__.py +@@ -0,0 +1,10 @@ ++# Code is adopted from tilelang/examples/deepseek_v32 ++# transformer/tilelang_kernel/__init__.py ++ ++from .sparse_mla_fwd import sparse_mla_fwd_interface ++from .sparse_mla_bwd import sparse_mla_bwd ++ ++__all__ = [ ++ "sparse_mla_fwd_interface", ++ "sparse_mla_bwd", ++] +\ No newline at end of file +diff --git a/megatron/core/transformer/tilelang_kernel/sparse_mla_bwd.py b/megatron/core/transformer/tilelang_kernel/sparse_mla_bwd.py +new file mode 100644 +index 000000000..83a259efa +--- /dev/null ++++ b/megatron/core/transformer/tilelang_kernel/sparse_mla_bwd.py +@@ -0,0 +1,272 @@ ++# ruff: noqa ++import tilelang ++from tilelang import language as T ++import torch ++ ++ ++@tilelang.jit(out_idx=[-1]) ++def preprocess( ++ B, ++ S, ++ H, ++ D, ++ block_ND=32, ++ num_stages=5, ++ dtype=T.bfloat16, ++ accum_dtype=T.float32, ++): ++ assert dtype == T.bfloat16 ++ assert accum_dtype == T.float32 ++ shape = [B, S, H, D] ++ ++ @T.prim_func ++ def preprocess_kernel( ++ O: T.Tensor(shape, dtype), ++ dO: T.Tensor(shape, dtype), ++ Delta: T.Tensor([B, S, H], accum_dtype), ++ ): ++ with T.Kernel(H, T.ceildiv(S, block_ND), B) as (bx, by, bz): ++ o = T.alloc_fragment([block_ND, block_ND], accum_dtype) ++ do = T.alloc_fragment([block_ND, block_ND], accum_dtype) ++ delta = T.alloc_fragment([block_ND], accum_dtype) ++ acc = T.alloc_fragment([block_ND, block_ND], accum_dtype) ++ T.clear(acc) ++ for k in T.Pipelined(T.ceildiv(D, block_ND), num_stages=num_stages): ++ T.copy(O[bz, by * block_ND : (by + 1) * block_ND, bx, k * block_ND : (k + 1) * block_ND], o) ++ T.copy(dO[bz, by * block_ND : (by + 1) * block_ND, bx, k * block_ND : (k + 1) * block_ND], do) ++ for i, j in T.Parallel(block_ND, block_ND): ++ acc[i, j] += o[i, j] * do[i, j] ++ T.reduce_sum(acc, delta, 1) ++ T.copy(delta, Delta[bz, by * block_ND : (by + 1) * block_ND, bx]) ++ ++ return preprocess_kernel ++ ++ ++@tilelang.jit(out_idx=[-1]) ++def postprocess( ++ B, ++ S_kv, ++ D, ++ D_tail, ++ kv_group=1, ++ block_N=64, ++ threads=256, ++ dtype=T.bfloat16, ++ accum_dtype=T.float32, ++): ++ assert dtype == T.bfloat16 ++ assert accum_dtype == T.float32 ++ dkv_shape = [B, S_kv, kv_group, D + D_tail] ++ ++ @T.prim_func ++ def postprocess_kernel( ++ dKV: T.Tensor(dkv_shape, accum_dtype), ++ dKV_out: T.Tensor(dkv_shape, dtype), ++ ): ++ with T.Kernel(T.ceildiv(S_kv, block_N), kv_group, B, threads=threads) as (bx, by, bz): ++ T.copy( ++ dKV[bz, bx * block_N : (bx + 1) * block_N, by, :], ++ dKV_out[bz, bx * block_N : (bx + 1) * block_N, by, :], ++ ) ++ ++ return postprocess_kernel ++ ++ ++@tilelang.jit( ++ out_idx=[-2], ++ pass_configs={ ++ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, ++ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, ++ tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: True, ++ }, ++) ++def bwd( ++ B, ++ S, ++ S_kv, ++ H, ++ D, ++ D_tail, ++ topk, ++ kv_group=1, ++ sm_scale=None, ++ is_causal=True, ++ block_size=32, ++ num_stages=0, ++ threads=128, ++ indices_dtype=T.int32, ++ dtype=T.bfloat16, ++ accum_dtype=T.float32, ++ masks_dtype=T.bool, ++): ++ assert is_causal == True, "non-casual is not supported now" ++ assert topk % block_size == 0, "otherwise will load some index=0 thus causing wrong kv to be loaded" ++ assert dtype == T.bfloat16 ++ assert accum_dtype == T.float32 ++ assert indices_dtype == T.int32 ++ ++ if sm_scale is None: ++ sm_scale = (D + D_tail) ** (-0.5) ++ sm_scale_mul_reciprocal_log2 = sm_scale * 1.44269504 # log2(e) ++ ++ H_kv = H // kv_group ++ q_shape = [B, S, H, D + D_tail] ++ k_shape = [B, S_kv, kv_group, D + D_tail] ++ o_shape = [B, S, H, D] ++ indices_shape = [B, S, kv_group, topk] ++ delta_shape = [B, S, H] ++ lse_shape = [B, S, H] ++ masks_shape = [B, S, kv_group, S_kv] ++ assert indices_dtype == T.int32 ++ assert dtype == T.bfloat16 ++ assert accum_dtype == T.float32 ++ ++ H = H_kv ++ padded_H = max(tilelang.math.next_power_of_2(H_kv), 16) ++ block_H = min(64, padded_H) ++ assert padded_H % block_H == 0 ++ NH = padded_H // block_H ++ BS = block_size ++ NS = tilelang.cdiv(topk, block_size) ++ ++ split_store = 2 ++ ++ @T.prim_func ++ def sparse_mla_bwd_kernel( ++ Q: T.Tensor(q_shape, dtype), ++ KV: T.Tensor(k_shape, dtype), ++ dO: T.Tensor(o_shape, dtype), ++ Indices: T.Tensor(indices_shape, indices_dtype), ++ Masks: T.Tensor(masks_shape, masks_dtype), ++ Lse: T.Tensor(lse_shape, accum_dtype), ++ Delta: T.Tensor(delta_shape, accum_dtype), ++ dQ: T.Tensor(q_shape, dtype), ++ dKV: T.Tensor(k_shape, accum_dtype), ++ ): ++ with T.Kernel(S, B, kv_group * NH, threads=threads) as (s_i, by, bz): ++ Q_shared = T.alloc_shared([block_H, D], dtype) ++ Q_tail_shared = T.alloc_shared([block_H, D_tail], dtype) ++ KV_shared = T.alloc_shared([BS, D], dtype) ++ KV_tail_shared = T.alloc_shared([BS, D_tail], dtype) ++ dO_shared = T.alloc_shared([block_H, D], dtype) ++ mask = T.alloc_fragment([BS], "bool") ++ ++ P_shared_cast = T.alloc_shared([block_H, BS], dtype) ++ dP_shared_cast = T.alloc_shared([block_H, BS], dtype) ++ dQ_shared = T.alloc_shared([block_H, D], dtype) ++ dQ_tail_shared = T.alloc_shared([block_H, D_tail], dtype) ++ ++ acc_p = T.alloc_fragment([block_H, BS], accum_dtype) ++ acc_dp = T.alloc_fragment([block_H, BS], accum_dtype) ++ acc_dq = T.alloc_fragment([block_H, D], accum_dtype) ++ acc_dq_tail = T.alloc_fragment([block_H, D_tail], accum_dtype) ++ acc_dkv = T.alloc_fragment([BS, D], accum_dtype) ++ acc_dkv_tail = T.alloc_fragment([BS, D_tail], accum_dtype) ++ acc_dkv_shared = T.alloc_shared([BS // split_store, D], accum_dtype) ++ acc_dkv_tail_shared = T.alloc_shared([BS // split_store, D_tail], accum_dtype) ++ ++ T.copy(Q[by, s_i, bz * block_H : (bz + 1) * block_H, :D], Q_shared) ++ T.copy(Q[by, s_i, bz * block_H : (bz + 1) * block_H, D:], Q_tail_shared) ++ T.copy(dO[by, s_i, bz * block_H : (bz + 1) * block_H, :D], dO_shared) ++ ++ T.clear(acc_dq) ++ T.clear(acc_dq_tail) ++ ++ # Process each block of indices ++ for i_i in T.Pipelined(NS, num_stages=num_stages): ++ # Compute attention scores ++ for bi_i in T.Parallel(BS): ++ mask[bi_i] = Masks[by, s_i, bz // NH, Indices[by, s_i, bz // NH, i_i * BS + bi_i]] ++ ++ for h_i, bi_i in T.Parallel(block_H, BS): ++ acc_p[h_i, bi_i] = T.if_then_else(mask[bi_i], -T.infinity(acc_p.dtype), 0) ++ ++ # Load KV, V for this block of indices ++ for bi_i, d_i in T.Parallel(BS, D): ++ KV_shared[bi_i, d_i] = KV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i], bz // NH, d_i] ++ ++ T.gemm(Q_shared, KV_shared, acc_p, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) ++ ++ for bi_i, d_i in T.Parallel(BS, D_tail): ++ KV_tail_shared[bi_i, d_i] = KV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i], bz // NH, D + d_i] ++ T.gemm(Q_tail_shared, KV_tail_shared[:, :D_tail], acc_p, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) ++ ++ for h_i, bi_i in T.Parallel(block_H, BS): ++ acc_p[h_i, bi_i] = T.exp2(acc_p[h_i, bi_i] * sm_scale_mul_reciprocal_log2 - Lse[by, s_i, bz * block_H + h_i]) ++ ++ T.copy(acc_p, P_shared_cast) ++ ++ T.gemm(dO_shared, KV_shared, acc_dp, transpose_B=True, policy=T.GemmWarpPolicy.FullCol, clear_accum=True) ++ ++ for h_i, bi_i in T.Parallel(block_H, BS): ++ acc_dp[h_i, bi_i] = acc_p[h_i, bi_i] * (acc_dp[h_i, bi_i] - Delta[by, s_i, bz * block_H + h_i]) * sm_scale ++ ++ T.copy(acc_dp, dP_shared_cast) ++ T.gemm(dP_shared_cast, KV_shared, acc_dq, policy=T.GemmWarpPolicy.FullCol) ++ T.gemm(dP_shared_cast, KV_tail_shared, acc_dq_tail, policy=T.GemmWarpPolicy.FullCol) ++ ++ T.gemm(dP_shared_cast, Q_shared, acc_dkv, transpose_A=True, policy=T.GemmWarpPolicy.FullCol, clear_accum=True) ++ T.gemm(P_shared_cast, dO_shared, acc_dkv, transpose_A=True, policy=T.GemmWarpPolicy.FullCol) ++ ++ T.clear(acc_dkv_tail) ++ T.gemm(dP_shared_cast, Q_tail_shared, acc_dkv_tail, transpose_A=True, policy=T.GemmWarpPolicy.FullCol) ++ ++ for s in range(split_store): ++ for bi_i, d_i in T.Parallel(BS, D): ++ if bi_i < BS // split_store: ++ acc_dkv_shared[bi_i, d_i] = acc_dkv[bi_i + s * (BS // split_store), d_i] ++ ++ for bi_i, d_i in T.Parallel(BS, D_tail): ++ if bi_i < BS // split_store: ++ acc_dkv_tail_shared[bi_i, d_i] = acc_dkv_tail[bi_i + s * (BS // split_store), d_i] ++ ++ for bi_i, d_i in T.Parallel(BS // split_store, D // 4): ++ T.atomic_addx4( ++ dKV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i + s * (BS // split_store)], bz // NH, d_i * 4], ++ acc_dkv_shared[bi_i, d_i * 4], ++ ) ++ ++ # Atomically update dKV, dKV_tail tensors ++ for bi_i, d_i in T.Parallel(BS // split_store, D_tail // 4): ++ T.atomic_addx4( ++ dKV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i + s * (BS // split_store)], bz // NH, D + d_i * 4], ++ acc_dkv_tail_shared[bi_i, d_i * 4], ++ ) ++ ++ # Store the accumulated dQ ++ T.copy(acc_dq, dQ_shared) ++ T.copy(acc_dq_tail[:, :D_tail], dQ_tail_shared) ++ ++ T.copy(dQ_shared, dQ[by, s_i, bz * block_H : (bz + 1) * block_H, :D]) ++ T.copy(dQ_tail_shared, dQ[by, s_i, bz * block_H : (bz + 1) * block_H, D:]) ++ ++ return sparse_mla_bwd_kernel ++ ++ ++def sparse_mla_bwd(q, kv, o, do, indices, masks, lse, dim_v, sm_scale=None, is_casual=True, return_kernel=False, delta=None): ++ assert q.is_contiguous() ++ assert kv.is_contiguous() ++ assert indices.is_contiguous() ++ assert lse.is_contiguous() ++ B, S, H, dim_plus_tail_dim = q.shape ++ _, S_kv, kv_group, _ = kv.shape ++ assert kv.shape[-1] == dim_plus_tail_dim ++ assert kv.shape[0] == B ++ # dim should be assigned ++ D = dim_v ++ ++ D_tail = dim_plus_tail_dim - D ++ topk = indices.shape[-1] ++ assert indices.shape == (B, S, kv_group, topk) ++ assert lse.shape == (B, S, H) ++ ++ # Get kernels ++ preprocess_kernel = preprocess(B, S, H, D) ++ bwd_kernel = bwd(B, S, S_kv, H, D, D_tail, topk, kv_group, sm_scale, is_casual) ++ ++ if delta is None: ++ delta = preprocess_kernel(o, do) ++ dkv = torch.zeros_like(kv, dtype=torch.float32) ++ dq = bwd_kernel(q, kv, do, indices, masks, lse, delta, dkv) ++ ++ return dq, dkv +\ No newline at end of file +diff --git a/megatron/core/transformer/tilelang_kernel/sparse_mla_fwd.py b/megatron/core/transformer/tilelang_kernel/sparse_mla_fwd.py +new file mode 100644 +index 000000000..e247038de +--- /dev/null ++++ b/megatron/core/transformer/tilelang_kernel/sparse_mla_fwd.py +@@ -0,0 +1,191 @@ ++# ruff: noqa ++import torch ++import tilelang ++from tilelang import language as T ++ ++ ++@tilelang.jit( ++ out_idx=[-2, -1], ++ pass_configs={ ++ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, ++ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, ++ }, ++) ++def sparse_mla_fwd( ++ heads, ++ dim, ++ tail_dim, ++ topk, ++ kv_group=1, ++ sm_scale=None, ++ is_causal=True, ++ CP0=True, ++ block_I=64, ++ num_stages=2, ++ threads=256, ++): ++ assert dim == tilelang.math.next_power_of_2(dim), f"haven't check padding correctness yet, dim={dim}" ++ assert tail_dim == tilelang.math.next_power_of_2(tail_dim), f"haven't check padding correctness yet, dim={tail_dim}" ++ assert is_causal == True, "non-casual is not supported" ++ assert topk % block_I == 0, "otherwise will load some index=0 thus causing wrong kv to be loaded" ++ if sm_scale is None: ++ sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) ++ else: ++ sm_scale = sm_scale * 1.44269504 # log2(e) ++ ++ batch = T.dynamic("batch") ++ seq_len = T.dynamic("seq_len") ++ seq_len_kv = T.dynamic("seq_len_kv") ++ ++ head_kv = heads // kv_group ++ q_shape = [batch, seq_len, heads, dim + tail_dim] ++ kv_shape = [batch, seq_len_kv, kv_group, dim + tail_dim] ++ o_shape = [batch, seq_len, heads, dim] ++ indices_shape = [batch, seq_len, kv_group, topk] ++ lse_shape = [batch, seq_len, heads] ++ masks_shape = [batch, seq_len, kv_group, seq_len_kv] ++ ++ masks_dtype = T.bool ++ indices_dtype = T.int32 ++ dtype = T.bfloat16 ++ accum_dtype = T.float32 ++ ++ G = kv_group ++ H = head_kv ++ padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) ++ if padded_H != H: ++ assert kv_group == 1, ( ++ "here we solve the H padding automatically, other wise you should handle Q copy and Output copy with your mask (when kv_group == 1, use g_i * padded_H:(g_i+1) * padded_H would be handled automatically)" ++ ) ++ BI = block_I ++ NI = tilelang.cdiv(topk, block_I) ++ D = dim ++ D_tail = tail_dim ++ ++ if head_kv > 64: ++ assert head_kv % 64 == 0, "head_kv should be a multiple of 64" ++ REPLICATE_H = head_kv // 64 ++ else: ++ REPLICATE_H = 1 ++ ++ H_per_block = padded_H if REPLICATE_H == 1 else 64 ++ ++ @T.prim_func ++ def main( ++ Q: T.Tensor(q_shape, dtype), # type: ignore ++ KV: T.Tensor(kv_shape, dtype), # type: ignore ++ Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore ++ Masks: T.Tensor(masks_shape, masks_dtype), # type: ignore ++ Output: T.Tensor(o_shape, dtype), # type: ignore ++ Lse: T.Tensor(lse_shape, accum_dtype), # type: ignore ++ ): ++ with T.Kernel(seq_len * REPLICATE_H, batch, kv_group, threads=threads) as ( ++ bx, ++ by, ++ bz, ++ ): ++ Q_shared = T.alloc_shared([H_per_block, D], dtype) ++ Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) ++ KV_shared = T.alloc_shared([BI, D], dtype) ++ K_tail_shared = T.alloc_shared([BI, D_tail], dtype) ++ O_shared = T.alloc_shared([H_per_block, D], dtype) ++ Lse_shared = T.alloc_shared([H_per_block], accum_dtype) ++ mask = T.alloc_fragment([BI], "bool") ++ ++ acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) ++ acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) ++ S_shared = T.alloc_shared([H_per_block, BI], dtype) ++ sumexp = T.alloc_fragment([H_per_block], accum_dtype) ++ sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) ++ alpha = T.alloc_fragment([H_per_block], accum_dtype) ++ m_i = T.alloc_fragment([H_per_block], accum_dtype) ++ m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) ++ ++ T.fill(acc_o, 0) ++ T.fill(sumexp, 0) ++ T.fill(m_i, -(2**30)) # avoid -inf - inf to cause nan ++ ++ b_i, g_i = by, bz ++ s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) ++ q_i = s_i ++ ++ H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64) ++ H1 = H0 + H_per_block ++ ++ T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) ++ T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) ++ ++ for i_i in T.Pipelined(NI, num_stages=num_stages): ++ for bi_i in T.Parallel(BI): ++ mask[bi_i] = Masks[b_i, s_i, g_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i]] ++ ++ for bi_i, d_i in T.Parallel(BI, D): ++ KV_shared[bi_i, d_i] = KV[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i] ++ for bi_i, d_i in T.Parallel(BI, D_tail): ++ K_tail_shared[bi_i, d_i] = KV[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, D + d_i] ++ for h_i, bi_i in T.Parallel(H_per_block, BI): ++ acc_s[h_i, bi_i] = T.if_then_else(mask[bi_i], -T.infinity(acc_s.dtype), 0) ++ T.gemm( ++ Q_shared, ++ KV_shared, ++ acc_s, ++ transpose_B=True, ++ policy=T.GemmWarpPolicy.FullRow, ++ ) ++ T.gemm( ++ Q_tail_shared, ++ K_tail_shared, ++ acc_s, ++ transpose_B=True, ++ policy=T.GemmWarpPolicy.FullRow, ++ ) ++ T.copy(m_i, m_i_prev) ++ T.reduce_max(acc_s, m_i, dim=1, clear=False) ++ for h_i in T.Parallel(H_per_block): ++ m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i]) ++ for h_i in T.Parallel(H_per_block): ++ alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) ++ for h_i, bi_i in T.Parallel(H_per_block, BI): ++ acc_s[h_i, bi_i] = T.exp2(acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale) ++ T.reduce_sum(acc_s, sumexp_i, dim=1) # is this a accumulate operator? ++ for h_i in T.Parallel(H_per_block): ++ sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] ++ for h_i, d_i in T.Parallel(H_per_block, D): ++ acc_o[h_i, d_i] = acc_o[h_i, d_i] * alpha[h_i] ++ ++ T.copy(acc_s, S_shared) ++ T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) ++ ++ # Rescale ++ for h_i, d_i in T.Parallel(H_per_block, D): ++ acc_o[h_i, d_i] /= sumexp[h_i] ++ for h_i in T.Parallel(H_per_block): ++ sumexp[h_i] = T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale ++ ++ T.copy(acc_o, O_shared) ++ T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) ++ T.copy(sumexp, Lse_shared) ++ T.copy(sumexp, Lse[b_i, s_i, H0:H1]) ++ ++ return main ++ ++ ++def sparse_mla_fwd_interface(q, kv, indices, masks, d_v, sm_scale=None, return_p_sum: bool = False, block_I=64, num_stages=2, threads=256): ++ is_casual = True ++ assert return_p_sum == False, "This kernel file is for fwd only" ++ assert q.is_contiguous() and kv.is_contiguous() and indices.is_contiguous() ++ batch, seq_len, heads, dim_plus_tail_dim = q.shape ++ _, seq_len_kv, kv_group, _ = kv.shape ++ ++ assert kv.shape[-1] == dim_plus_tail_dim ++ tail_dim = dim_plus_tail_dim - d_v ++ assert kv.shape[0] == batch ++ _, _, _, topk = indices.shape ++ assert indices.shape == (batch, seq_len, kv_group, topk) ++ assert masks.shape == (batch, seq_len, kv_group, seq_len_kv) ++ ++ kernel = sparse_mla_fwd( ++ heads, d_v, tail_dim, topk, kv_group, sm_scale, is_casual, block_I=block_I, num_stages=num_stages, threads=threads ++ ) ++ out, lse = kernel(q, kv, indices, masks) ++ return out, lse +\ No newline at end of file +diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py +index e2705bd9f..29a0ff9e0 100644 +--- a/megatron/core/transformer/transformer_config.py ++++ b/megatron/core/transformer/transformer_config.py +@@ -935,11 +935,10 @@ class TransformerConfig(ModelParallelConfig): + f" but got {self.context_parallel_size=}." + ) + elif self.experimental_attention_variant == "dsa": +- assert ( +- self.context_parallel_size == 1 +- ), "Currently context parallelism is not supported by DSAttention!" ++ # assert ( ++ # self.context_parallel_size == 1 ++ # ), "Currently context parallelism is not supported by DSAttention!" + assert not self.apply_rope_fusion, "RoPE fusion is not supported for DSAttention" +- + if self.fp8: + # cannot support first last layer bf16 with delayed scaling + if self.first_last_layers_bf16 and self.fp8_recipe == Fp8Recipe.delayed: diff --git a/docker/deepseekv32/transformers.patch b/docker/deepseekv32/transformers.patch new file mode 100644 index 00000000000..a7631aa00c2 --- /dev/null +++ b/docker/deepseekv32/transformers.patch @@ -0,0 +1,20 @@ +diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py +index f6a12e7cef..22129a86ee 100644 +--- a/src/transformers/models/auto/configuration_auto.py ++++ b/src/transformers/models/auto/configuration_auto.py +@@ -1355,6 +1355,15 @@ class AutoConfig: + "Detected mistral model with layer_types, treating as ministral for alternating attention compatibility. " + ) + config_dict["model_type"] = "ministral" ++ if config_dict["model_type"] == "deepseek_v32": ++ logger.info( ++ "Detected deepseek_v32 model, treating as deepseek_v3 for compatibility." ++ ) ++ config_dict["model_type"] = "deepseek_v3" ++ if "architectures" in config_dict: ++ config_dict["architectures"] = [ ++ arch.replace("DeepseekV32", "DeepseekV3") for arch in config_dict["architectures"] ++ ] + + try: + config_class = CONFIG_MAPPING[config_dict["model_type"]] diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index a92198a6744..4bb7548ea75 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -462,6 +462,9 @@ def update_weights(self) -> None: if self.args.offload_train: reload_process_groups() + if isinstance(num_new_engines, tuple): + num_new_engines = num_new_engines[0] + if num_new_engines > 0: self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock) dist.barrier(group=get_gloo_group()) diff --git a/miles/backends/megatron_utils/megatron_to_hf/__init__.py b/miles/backends/megatron_utils/megatron_to_hf/__init__.py index 84ff899aa52..b9b394cbcff 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/__init__.py +++ b/miles/backends/megatron_utils/megatron_to_hf/__init__.py @@ -1,4 +1,5 @@ from .deepseekv3 import convert_deepseekv3_to_hf +from .deepseekv32 import convert_deepseekv32_to_hf from .glm4 import convert_glm4_to_hf from .glm4moe import convert_glm4moe_to_hf from .llama import convert_llama_to_hf @@ -41,6 +42,8 @@ def _convert_to_hf_core(args, model_name, name, param): converted_named_tensors = convert_qwen3_next_to_hf(args, name, param) elif "qwen2" in model_name or "qwen3" in model_name: converted_named_tensors = convert_qwen2_to_hf(args, name, param) + elif "deepseekv32" in model_name: + converted_named_tensors = convert_deepseekv32_to_hf(args, name, param) elif "deepseekv3" in model_name: converted_named_tensors = convert_deepseekv3_to_hf(args, name, param) diff --git a/miles/backends/megatron_utils/megatron_to_hf/deepseekv32.py b/miles/backends/megatron_utils/megatron_to_hf/deepseekv32.py new file mode 100644 index 00000000000..3b6519ada8e --- /dev/null +++ b/miles/backends/megatron_utils/megatron_to_hf/deepseekv32.py @@ -0,0 +1,135 @@ +import re + +import sglang +import torch +from packaging.version import parse + + +def convert_deepseekv32_to_hf(args, name, param): + if name == "module.module.embedding.word_embeddings.weight": + return [("model.embed_tokens.weight", param)] + if name == "module.module.output_layer.weight": + return [("lm_head.weight", param)] + if name == "module.module.decoder.final_layernorm.weight": + return [("model.norm.weight", param)] + + try: + head_dim = args.kv_channels if args.kv_channels is not None else args.hidden_size // args.num_attention_heads + except AttributeError: + head_dim = args.hidden_size // args.num_attention_heads + value_num_per_group = args.num_attention_heads // args.num_query_groups + + decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)" + match = re.match(decoder_layers_pattern, name) + if match: + layer_idx, rest = match.groups() + + # experts + expert_pattern = r"mlp.experts\.(.+)\.weight(\d+)" + match = re.match(expert_pattern, rest) + if match: + rest, expert_idx = match.groups() + if rest == "linear_fc1": + gate_weight, up_weight = param.chunk(2, dim=0) + outputs = [ + (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight", gate_weight), + (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight", up_weight), + ] + return outputs + elif rest == "linear_fc2": + outputs = [ + (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight", param), + ] + if parse(sglang.__version__) < parse("0.4.9.post5") and args.sglang_enable_ep_moe: + outputs += [ + ( + f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.input_scale", + torch.tensor(1.0, dtype=torch.float32, device=param.device), + ), + ( + f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight_scale", + torch.tensor(1.0, dtype=torch.float32, device=param.device), + ), + ] + return outputs + else: + raise ValueError(f"Unknown expert parameter name: {name}") + + # shared expert + shared_expert_pattern = r"mlp.shared_experts\.(.+)" + match = re.match(shared_expert_pattern, rest) + if match: + rest = match.groups()[0] + if rest == "linear_fc1.weight": + gate_weight, up_weight = param.chunk(2, dim=0) + return [ + (f"model.layers.{layer_idx}.mlp.shared_experts.gate_proj.weight", gate_weight), + (f"model.layers.{layer_idx}.mlp.shared_experts.up_proj.weight", up_weight), + ] + elif rest == "linear_fc2.weight": + return [(f"model.layers.{layer_idx}.mlp.shared_experts.down_proj.weight", param)] + else: + raise ValueError(f"Unknown shared expert parameter name: {name}") + + if rest == "self_attention.linear_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.o_proj.weight", param)] + elif rest == "self_attention.linear_q_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.q_proj.weight", param)] + elif rest == "self_attention.linear_q_down_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.q_a_proj.weight", param)] + elif rest == "self_attention.q_layernorm.weight": + return [(f"model.layers.{layer_idx}.self_attn.q_a_layernorm.weight", param)] + elif rest == "self_attention.linear_q_up_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.q_b_proj.weight", param)] + elif rest == "self_attention.linear_qkv.bias": + param = param.view(args.num_query_groups, -1) + q_bias, k_bias, v_bias = torch.split( + param, + split_size_or_sections=[value_num_per_group * head_dim, head_dim, head_dim], + dim=1, + ) + q_bias = q_bias.contiguous().flatten() + k_bias = k_bias.contiguous().flatten() + v_bias = v_bias.contiguous().flatten() + return [ + (f"model.layers.{layer_idx}.self_attn.q_proj.bias", q_bias), + (f"model.layers.{layer_idx}.self_attn.k_proj.bias", k_bias), + (f"model.layers.{layer_idx}.self_attn.v_proj.bias", v_bias), + ] + elif rest == "mlp.linear_fc1.weight": + gate_weight, up_weight = param.chunk(2, dim=0) + return [ + (f"model.layers.{layer_idx}.mlp.gate_proj.weight", gate_weight), + (f"model.layers.{layer_idx}.mlp.up_proj.weight", up_weight), + ] + elif rest == "mlp.linear_fc2.weight": + return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)] + elif rest == "self_attention.linear_qkv.layer_norm_weight" or rest == "input_layernorm.weight": + return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)] + elif rest == "mlp.linear_fc1.layer_norm_weight": + return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)] + elif rest == "self_attention.linear_kv_down_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.kv_a_proj_with_mqa.weight", param)] + elif rest == "self_attention.kv_layernorm.weight": + return [(f"model.layers.{layer_idx}.self_attn.kv_a_layernorm.weight", param)] + elif rest == "self_attention.linear_kv_up_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.kv_b_proj.weight", param)] + elif rest == "pre_mlp_layernorm.weight": + return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)] + # DSA Indexer parameters + elif rest == "self_attention.core_attention.indexer.linear_wq_b.weight": + return [(f"model.layers.{layer_idx}.self_attn.indexer.wq_b.weight", param)] + elif rest == "self_attention.core_attention.indexer.linear_wk.weight": + return [(f"model.layers.{layer_idx}.self_attn.indexer.wk.weight", param)] + elif rest == "self_attention.core_attention.indexer.k_norm.weight": + return [(f"model.layers.{layer_idx}.self_attn.indexer.k_norm.weight", param)] + elif rest == "self_attention.core_attention.indexer.k_norm.bias": + return [(f"model.layers.{layer_idx}.self_attn.indexer.k_norm.bias", param)] + elif rest == "self_attention.core_attention.indexer.linear_weights_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.indexer.weights_proj.weight", param)] + elif rest == "mlp.router.weight": + return [(f"model.layers.{layer_idx}.mlp.gate.weight", param)] + elif rest == "mlp.router.expert_bias": + return [(f"model.layers.{layer_idx}.mlp.gate.e_score_correction_bias", param)] + + raise ValueError(f"Unknown parameter name: {name}") diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py index 41495f3969b..da5b2b55ad5 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py @@ -42,7 +42,10 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio # TODO: find a clearer way. if converted_name.endswith("_scale"): continue - quantize_named_params.extend(_quantize_param(converted_name, param, weight_block_size)) + if_use_ue8m0_in_moe = True if args.sglang_moe_runner_backend == "deep_gemm" else False + quantize_named_params.extend( + _quantize_param(converted_name, param, weight_block_size, if_use_ue8m0_in_moe=if_use_ue8m0_in_moe) + ) return quantize_named_params @@ -72,6 +75,9 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio "self_attention.linear_q_up_proj.weight", "self_attention.linear_kv_down_proj.weight", "self_attention.linear_kv_up_proj.weight", + # dsa indexer + "self_attention.core_attention.indexer.linear_wq_b.weight", + "self_attention.core_attention.indexer.linear_wk.weight", ]: quantize_named_params = [] for converted_name, param in converted_named_params: @@ -83,13 +89,15 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio return converted_named_params -def _quantize_param(name, weight, weight_block_size): +def _quantize_param(name, weight, weight_block_size, if_use_ue8m0_in_moe=True): assert name.endswith(".weight"), f"Expected weight parameter, got {name}" FP8_MIN = torch.finfo(torch.float8_e4m3fn).min FP8_MAX = torch.finfo(torch.float8_e4m3fn).max if weight_block_size is not None: - if should_deepgemm_weight_requant_ue8m0 and should_deepgemm_weight_requant_ue8m0( - weight_block_size=weight_block_size + if ( + should_deepgemm_weight_requant_ue8m0 + and should_deepgemm_weight_requant_ue8m0(weight_block_size=weight_block_size) + and if_use_ue8m0_in_moe ): qweight, scale = quant_weight_ue8m0(weight, weight_block_size=weight_block_size) scale = transform_scale_ue8m0(scale, mn=qweight.shape[-2]) diff --git a/miles/backends/megatron_utils/update_weight/common.py b/miles/backends/megatron_utils/update_weight/common.py index 85fe76a1b8d..e958566dcaa 100644 --- a/miles/backends/megatron_utils/update_weight/common.py +++ b/miles/backends/megatron_utils/update_weight/common.py @@ -203,6 +203,17 @@ def _named_params_and_buffers_global( yield f"module.module.mtp.layers.{layer_idx}.transformer_layer.mlp.experts.{rest}.weight{expert_idx}", param continue + # TODO: a hacking here, need to be cleaner + duplicated = [ + "indexer.linear_weights_proj", + "indexer.linear_wk", + "indexer.linear_wq_b", + "linear_q_down_proj", + "linear_kv_down_proj", + ] + if any(dup in name for dup in duplicated): + param.parallel_mode = "duplicated" + layer_idx, rest = match.groups() layer_idx = int(layer_idx) + layer_offset diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py index a7f88d13762..abc790761d0 100644 --- a/miles/backends/training_utils/loss.py +++ b/miles/backends/training_utils/loss.py @@ -858,7 +858,7 @@ def loss_function( return ( loss, - torch.tensor(num_tokens if args.calculate_per_token_loss else 1, device=logits.device), + torch.tensor(num_tokens if args.calculate_per_token_loss else 1, dtype=torch.int, device=logits.device), { "keys": list(log.keys()), "values": torch.tensor( diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 07102029247..f3dce7fbfe5 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1723,6 +1723,16 @@ def miles_validate_args(args): args.use_dynamic_batch_size is False ), "Dynamic batch size is not supported for bshd format. Please specify --micro-batch-size instead." + assert args.qkv_format in [ + "thd", + "bshd", + ], f"qkv_format {args.qkv_format} is not supported. (only 'thd' and 'bshd' are supported)" + if args.qkv_format == "bshd": + assert args.train_backend == "megatron", "bshd format is only supported for megatron backend." + assert ( + args.use_dynamic_batch_size is False + ), "Dynamic batch size is not supported for bshd format. Please specify --micro-batch-size instead." + def hf_validate_args(args, hf_config): def equal(x, y): diff --git a/miles/utils/data.py b/miles/utils/data.py index 6e64ef678de..d258b4070c1 100644 --- a/miles/utils/data.py +++ b/miles/utils/data.py @@ -200,13 +200,22 @@ def __init__( metadata["tools"] = tools if apply_chat_template: - output_prompt = tokenizer.apply_chat_template( - prompt, - tools=tools, - tokenize=False, - add_generation_prompt=True, - **(apply_chat_template_kwargs or {}), - ) + ### DSV32 + try: + prompt = tokenizer.apply_chat_template( + prompt, + tools, + tokenize=False, + add_generation_prompt=True, + **apply_chat_template_kwargs, + ) + except Exception: + from sglang.srt.entrypoints.openai.encoding_dsv32 import encode_messages + + encode_config = dict(thinking_mode="thinking", drop_thinking=True, add_default_bos_token=True) + prompt = encode_messages(prompt, **encode_config) + ### DSV32 + output_prompt = prompt else: output_prompt = prompt diff --git a/miles/utils/external_utils/command_utils.py b/miles/utils/external_utils/command_utils.py index 8c7c9316b95..bdfc864b425 100644 --- a/miles/utils/external_utils/command_utils.py +++ b/miles/utils/external_utils/command_utils.py @@ -26,6 +26,7 @@ def convert_checkpoint( extra_args: str = "", dir_dst: str = "/root", hf_checkpoint: str | None = None, + megatron_path: str = "/host_home/primary_synced/Megatron-LM", ): hf_checkpoint = hf_checkpoint or f"/root/models/{model_name}" @@ -51,14 +52,15 @@ def convert_checkpoint( exec_command( f"source {repo_base_dir}/scripts/models/{megatron_model_type}.sh && " - f"PYTHONPATH=/root/Megatron-LM " + # Use installed Megatron instead of hardcoded path + f"PYTHONPATH={megatron_path} " f"torchrun " f"--nproc-per-node {num_gpus_per_node} " f"{multinode_args}" f"tools/convert_hf_to_torch_dist.py " "${MODEL_ARGS[@]} " f"--hf-checkpoint {hf_checkpoint} " - f"--save {path_dst}" + f"--save {path_dst} " f"{extra_args}" ) @@ -67,9 +69,9 @@ def rsync_simple(path_src: str, path_dst: str): exec_command(f"mkdir -p {path_dst} && rsync -a --info=progress2 {path_src}/ {path_dst}") -def hf_download_dataset(full_name: str): +def hf_download_dataset(full_name: str, data_dir: str = "/root/datasets"): _, partial_name = full_name.split("/") - exec_command(f"hf download --repo-type dataset {full_name} --local-dir /root/datasets/{partial_name}") + exec_command(f"hf download --repo-type dataset {full_name} --local-dir {data_dir}/{partial_name}") def fp8_cast_bf16(path_src, path_dst): @@ -98,6 +100,7 @@ def execute_train( before_ray_job_submit=None, extra_env_vars=None, config: ExecuteTrainConfig | None = None, + megatron_path: str = "/host_home/primary_synced/Megatron-LM", ): if extra_env_vars is None: extra_env_vars = {} @@ -139,7 +142,8 @@ def execute_train( runtime_env_json = json.dumps( { "env_vars": { - "PYTHONPATH": "/root/Megatron-LM/", + # Use installed Megatron instead of hardcoded path + "PYTHONPATH": f"{megatron_path}", # If setting this in FSDP, the computation communication overlapping may have issues **( {} diff --git a/miles_plugins/mbridge/__init__.py b/miles_plugins/mbridge/__init__.py index f97c7f46eef..67b824aa94e 100644 --- a/miles_plugins/mbridge/__init__.py +++ b/miles_plugins/mbridge/__init__.py @@ -1,6 +1,29 @@ +import os +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../..")) + +from .deepseekv32 import DeepseekV32Bridge from .glm4 import GLM4Bridge from .glm4moe import GLM4MoEBridge from .mimo import MimoBridge from .qwen3_next import Qwen3NextBridge -__all__ = ["GLM4Bridge", "GLM4MoEBridge", "Qwen3NextBridge", "MimoBridge"] +__all__ = ["DeepseekV32Bridge", "GLM4Bridge", "GLM4MoEBridge", "Qwen3NextBridge", "MimoBridge"] + +from mbridge import AutoBridge + +_original_from_config = AutoBridge.from_config + + +@classmethod +def _patched_from_config(cls, hf_config, **kwargs): + if hasattr(hf_config, "index_n_heads"): + from mbridge.core.bridge import _MODEL_REGISTRY + + return _MODEL_REGISTRY["deepseek_v32"](hf_config, **kwargs) + + return _original_from_config(hf_config, **kwargs) + + +AutoBridge.from_config = _patched_from_config diff --git a/miles_plugins/mbridge/deepseekv32.py b/miles_plugins/mbridge/deepseekv32.py new file mode 100644 index 00000000000..19e417fa06c --- /dev/null +++ b/miles_plugins/mbridge/deepseekv32.py @@ -0,0 +1,57 @@ +from megatron.core.transformer.enums import AttnBackend + +from mbridge.core import register_model +from mbridge.models import DeepseekV3Bridge + + +@register_model("deepseek_v32") +class DeepseekV32Bridge(DeepseekV3Bridge): + + # Weights with parallel_mode="duplicated" that should NOT be gathered across TP + _DUPLICATED_WEIGHTS = { + "self_attention.core_attention.indexer.linear_wq_b.weight", + "self_attention.core_attention.indexer.linear_wk.weight", + "self_attention.core_attention.indexer.linear_weights_proj.weight", + } + + _ATTENTION_MAPPING = DeepseekV3Bridge._ATTENTION_MAPPING.copy() + + # Because the indexer needs the norm output, we cannot use the fused transformer engine impl and have to compute it separately. + if "self_attention.linear_q_up_proj.layer_norm_weight" in _ATTENTION_MAPPING: + del _ATTENTION_MAPPING["self_attention.linear_q_up_proj.layer_norm_weight"] + if "self_attention.linear_kv_up_proj.layer_norm_weight" in _ATTENTION_MAPPING: + del _ATTENTION_MAPPING["self_attention.linear_kv_up_proj.layer_norm_weight"] + + _ATTENTION_MAPPING.update( + { + "self_attention.q_layernorm.weight": ["model.layers.{layer_number}.self_attn.q_a_layernorm.weight"], + "self_attention.kv_layernorm.weight": ["model.layers.{layer_number}.self_attn.kv_a_layernorm.weight"], + "self_attention.core_attention.indexer.linear_wq_b.weight": [ + "model.layers.{layer_number}.self_attn.indexer.wq_b.weight" + ], + "self_attention.core_attention.indexer.linear_wk.weight": [ + "model.layers.{layer_number}.self_attn.indexer.wk.weight" + ], + "self_attention.core_attention.indexer.k_norm.weight": [ + "model.layers.{layer_number}.self_attn.indexer.k_norm.weight" + ], + "self_attention.core_attention.indexer.k_norm.bias": [ + "model.layers.{layer_number}.self_attn.indexer.k_norm.bias" + ], + "self_attention.core_attention.indexer.linear_weights_proj.weight": [ + "model.layers.{layer_number}.self_attn.indexer.weights_proj.weight" + ], + } + ) + + def _build_config(self): + config = super()._build_config() + + config.attention_backend = AttnBackend.auto + + config.experimental_attention_variant = "dsa" + config.dsa_indexer_n_heads = getattr(self.hf_config, "dsa_indexer_n_heads", 64) + config.dsa_indexer_head_dim = getattr(self.hf_config, "dsa_indexer_head_dim", 128) + config.dsa_indexer_topk = getattr(self.hf_config, "dsa_indexer_topk", 2048) + + return config diff --git a/scripts/models/deepseek-v32-5layer.sh b/scripts/models/deepseek-v32-5layer.sh new file mode 100644 index 00000000000..2466640afd5 --- /dev/null +++ b/scripts/models/deepseek-v32-5layer.sh @@ -0,0 +1 @@ +MODEL_ARGS_NUM_LAYERS=5 source "$(dirname -- "${BASH_SOURCE[0]}")/deepseek-v32.sh" diff --git a/scripts/models/deepseek-v32.sh b/scripts/models/deepseek-v32.sh new file mode 100644 index 00000000000..a98f2f561d1 --- /dev/null +++ b/scripts/models/deepseek-v32.sh @@ -0,0 +1,69 @@ +NLAYERS="${MODEL_ARGS_NUM_LAYERS:-61}" +FIRST_K_DENSE_REPLACE=3 + +arr=() +for ((i=0; i 1 else False, + extra_args=extra_args, + dir_dst=f"{args.model_dir}", + megatron_path=args.megatron_path, + ) + + +@app.command() +@U.dataclass_cli +def prepare_cp(args: ScriptArgs): + _prepare_cp(args) + + +def _prepare_cp(args: ScriptArgs): + U.rsync_simple( + path_src=f"{args.model_dir}/{args.model_name}_torch_dist", + path_dst=f"{args.model_local_dir}/{args.model_name}_torch_dist", + ) + U.rsync_simple( + path_src=f"{args.model_dir}/{args.model_name}", + path_dst=f"{args.model_local_dir}/{args.model_name}", + ) + + +@app.command() +@U.dataclass_cli +def train(args: ScriptArgs): + print("running on {args.num_nodes} nodes") + # ensure files are there is it was not synced before + # _prepare_cp(args) + + load_save_path = f"{args.save_dir}/{args.run_id}/checkpoints" + ckpt_args = ( + f"--hf-checkpoint {args.model_local_dir}/{args.model_name} " + f"--ref-load {args.model_local_dir}/{args.model_name}_torch_dist " + f"--load {load_save_path} " + f"--save {load_save_path} " + "--save-interval 20 " + "--save-retain-interval 20 " + ) + + rollout_args = ( + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type math " + "--num-rollout 3000 " + "--rollout-batch-size 1 " + "--n-samples-per-prompt 1 " + "--rollout-temperature 0.8 " + # ------------ + "--num-steps-per-rollout 1 " + "--balance-data " + ) + + if args.mode != "debug_minimal": + rollout_args += ( + "--over-sampling-batch-size 256 " + "--dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std " + ) + + # sometimes disable eval to speed up debugging + eval_args = "" + if (args.mode != "debug_minimal") and args.enable_eval: + eval_args += "--eval-interval 20 " "--eval-top-p 0.7 " + + match args.task: + case "dapo_aime": + rollout_args += ( + f"--prompt-data {args.data_dir}/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + f"--rollout-max-response-len {100 if args.mode == 'debug_minimal' else 8192} " + ) + eval_args += ( + f"--eval-prompt-data aime {args.data_dir}/aime-2024/aime-2024.jsonl " + "--n-samples-per-eval-prompt 8 " + "--eval-max-response-len 8192 " + ) + case "gsm8k": + rollout_args += ( + f"--prompt-data {args.data_dir}/gsm8k/train.parquet " + "--input-key messages " + # Deliberately make it very short for this easy task + "--rollout-max-response-len 256 " + ) + eval_args += ( + f"--eval-prompt-data gsm8k {args.data_dir}/gsm8k/test.parquet " + "--n-samples-per-eval-prompt 1 " + "--eval-max-response-len 256 " + ) + + if args.num_nodes <= 2: + perf_args = ( + "--tensor-model-parallel-size 2 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 2 " + "--expert-model-parallel-size 4 " + "--expert-tensor-parallel-size 1 " + ) + elif args.num_nodes <= 4: + # TODO remove this temp cfg + perf_args = ( + "--tensor-model-parallel-size 4 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 4 " + "--expert-model-parallel-size 4 " + "--expert-tensor-parallel-size 1 " + ) + else: + # TODO choose a good config (currently randomly change to suit 64gpu) + perf_args = ( + "--tensor-model-parallel-size 8 " + "--sequence-parallel " + f"--pipeline-model-parallel-size {1 if args.model_name == 'DeepSeek-V3.2-5layer' else 4} " + "--context-parallel-size 2 " + "--expert-model-parallel-size 16 " + "--expert-tensor-parallel-size 1 " + ) + if re.search(r"(\d+)layer", args.model_name) is None: + perf_args += "--decoder-last-pipeline-num-layers 13 " + perf_args += ( + # ------------ + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + # ------------ + # "--use-dynamic-batch-size " + "--micro-batch-size 1 " + # TODO temp use tiny value + "--max-tokens-per-gpu 2048 " + # "--max-tokens-per-gpu 16384 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + # TODO run-deepseek-r1.sh enables use-kl-loss but w/ coef 0. can we just disable it like this? + # "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + "--use-miles-router " + "--use-rollout-routing-replay " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + # ------------ + # "--optimizer-cpu-offload " + # "--overlap-cpu-optimizer-d2h-h2d " + # "--use-precision-aware-optimizer " + ) + + sglang_decode_max_bs = 256 + sglang_world_size = 4 if args.num_nodes <= 4 else 64 + sglang_attn_dp_size = 1 if args.num_nodes <= 4 else 8 + sglang_attn_tp_size = sglang_world_size // sglang_attn_dp_size + sglang_args = ( + f"--rollout-num-gpus-per-engine {sglang_world_size} " + "--sglang-mem-fraction-static 0.7 " + f"--sglang-tp-size {sglang_world_size} " + f"--sglang-ep-size {sglang_world_size} " + # dp attention + "--sglang-enable-dp-attention " + f"--sglang-dp-size {sglang_attn_dp_size} " + "--sglang-moe-dense-tp-size 1 " + "--sglang-enable-dp-lm-head " + # make every dp rank has 128 concurrency + "--sglang-server-concurrency 1024 " + f"--sglang-max-running-requests {sglang_world_size * sglang_decode_max_bs // sglang_attn_tp_size} " + f"--sglang-chunked-prefill-size {sglang_world_size * sglang_decode_max_bs} " + f"--sglang-cuda-graph-max-bs {sglang_decode_max_bs} " + "--sglang-disable-cuda-graph " + # For quick experiments + # """--sglang-json-model-override-args '{"num_hidden_layers": 5}' """ + ) + sglang_extra_env_vars = {} + + if args.enable_deepep: + sglang_args += ( + "--sglang-moe-a2a-backend deepep " + "--sglang-moe-runner-backend deep_gemm " + "--sglang-deepep-mode low_latency " + ) + sglang_extra_env_vars["SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = f"{sglang_decode_max_bs}" + + misc_args = ( + # default dropout in megatron is 0.1 + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + # should be good for model performance + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + # need to comment this when using model with MLA + # "--attention-backend flash " + f"--update-weight-buffer-size {4 * 1024 ** 3} " + # TODO maybe enable it + # use deepep for megatron + # "--moe-enable-deepep " + # "--moe-token-dispatcher-type flex " + # ------------ + f"--actor-num-nodes {args.num_nodes} " + f"--actor-num-gpus-per-node {args.num_gpus_per_node} " + f"--num-gpus-per-node {args.num_gpus_per_node} " + "--colocate " + "--use-fault-tolerance " + f"--dump-details /root/shared_data/{args.run_id}/dump_details " + "--disable-weights-backuper " + "--model-name deepseekv32 " # for mbridge load + "--train-memory-margin-bytes 1073741824 " + # "--check-weight-update-equal " + "--qkv-format bshd " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__, run_id=args.run_id)} " + f"{perf_args} " + f"{eval_args} " + f"{sglang_args} " + f"{misc_args} " + f"{args.extra_args} " + ) + + U.execute_train( + train_args=train_args, + config=args, + # TODO may get it from `config` + num_gpus_per_node=args.num_gpus_per_node, + megatron_model_type=args.megatron_model_type, + extra_env_vars={**sglang_extra_env_vars}, + megatron_path=args.megatron_path, + ) + + +if __name__ == "__main__": + app()