diff --git a/docker/Dockerfile b/docker/Dockerfile index d145b6a73..6b5c553d4 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -45,6 +45,17 @@ RUN git clone https://github.com/Dao-AILab/flash-attention.git && \ RUN pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c --no-deps RUN pip install flash-linear-attention==0.4.1 +# FlashQLA: optional GDN backend for Qwen3.5/Qwen3-Next (--qwen-gdn-backend flashqla; requires SM90+). +# Mirrors slime's split (x86 docker/Dockerfile installs it; arm Dockerfile.gb10 makes it opt-in): +# vime has one Dockerfile, so gate on INSTALL_FLASHQLA (justfile sets =1 for the x86 release, +# leaves =0 for arm). Verified on gb200/Blackwell: the wheel builds but `import flash_qla` raises +# "FlashQLA now support sm90 only" and the install downgrades tilelang (breaks vLLM) — so arm skips. +ARG INSTALL_FLASHQLA=0 +RUN if [ "${INSTALL_FLASHQLA}" = "1" ]; then \ + pip install git+https://github.com/QwenLM/FlashQLA.git --no-build-isolation; \ + else \ + echo "Skipping FlashQLA (INSTALL_FLASHQLA=0; sm90/Hopper-only — use --qwen-gdn-backend fla)"; \ + fi RUN pip install tilelang -f https://tile-ai.github.io/whl/nightly/cu128/ # cublas + (cu13) cuda dev headers. The arm64 (sbsa) vllm/vllm-openai base ships diff --git a/docker/justfile b/docker/justfile index 289f903c7..3e67029a8 100644 --- a/docker/justfile +++ b/docker/justfile @@ -24,7 +24,7 @@ BUILDER := "vime-builder" # Default — cu12.9, no cu marker in the tag. build: - ARG_TAG_SUFFIX="" ARG_BUILD_EXTRA_ARGS="" just _build-digest + ARG_TAG_SUFFIX="" ARG_BUILD_EXTRA_ARGS="--build-arg INSTALL_FLASHQLA=1" just _build-digest # cu13 variant — vLLM default-CUDA base; ENABLE_CUDA_13 builds the CUDA-13 # TransformerEngine/Triton on top. @@ -71,7 +71,7 @@ build-test: cd .. VERSION="$(cat docker/version.txt | tr -d '\n')" - docker build -f docker/Dockerfile . --build-arg HTTP_PROXY="$http_proxy" --build-arg HTTPS_PROXY="$https_proxy" --build-arg NO_PROXY="localhost,127.0.0.1" -t "{{IMAGE}}:vime-test-${VERSION}" + docker build -f docker/Dockerfile . --build-arg HTTP_PROXY="$http_proxy" --build-arg HTTPS_PROXY="$https_proxy" --build-arg NO_PROXY="localhost,127.0.0.1" --build-arg INSTALL_FLASHQLA=1 -t "{{IMAGE}}:vime-test-${VERSION}" docker push "{{IMAGE}}:vime-test-${VERSION}" docker tag "{{IMAGE}}:vime-test-${VERSION}" "{{IMAGE}}:vime-test-latest" diff --git a/docker/patch/latest/megatron.patch b/docker/patch/latest/megatron.patch index 189ba8164..3be8152b8 100644 --- a/docker/patch/latest/megatron.patch +++ b/docker/patch/latest/megatron.patch @@ -48,98 +48,146 @@ index a5b6c009b..22794d7e6 100644 + allow_partial_load=True, ), ) - + diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py -index 55179ff30..43ef12a29 100644 +index 55179ff30..6629f41a6 100644 --- a/megatron/core/distributed/distributed_data_parallel.py +++ b/megatron/core/distributed/distributed_data_parallel.py -@@ -45,6 +45,7 @@ class DistributedDataParallel(_BaseDataParallel): +@@ -45,6 +45,8 @@ class DistributedDataParallel(_BaseDataParallel): module: torch.nn.Module, disable_bucketing: bool = False, pg_collection: Optional[ProcessGroupCollection] = None, + disable_grad_buffers_cpu_backup: bool = False, ++ disable_param_buffers_cpu_backup: bool = False, ): super().__init__(config=config, module=module) if has_config_logger_enabled(config): -@@ -209,6 +210,7 @@ class DistributedDataParallel(_BaseDataParallel): +@@ -209,6 +211,8 @@ class DistributedDataParallel(_BaseDataParallel): param_and_grad_dtype_to_indices[(param_dtype, grad_dtype)], self.ddp_config.nccl_ub, pg_collection, + disable_grad_buffers_cpu_backup=disable_grad_buffers_cpu_backup, ++ disable_param_buffers_cpu_backup=disable_param_buffers_cpu_backup, ) ) - + diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py -index 088374fbf..7a416f8e4 100644 +index 088374fbf..a9982e176 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py -@@ -599,6 +599,7 @@ class _ParamAndGradBuffer: +@@ -599,6 +599,8 @@ class _ParamAndGradBuffer: param_indices: List[int], nccl_ub: bool, pg_collection: Optional[ProcessGroupCollection] = None, + disable_grad_buffers_cpu_backup: bool = False, ++ disable_param_buffers_cpu_backup: bool = False, ): - + if pg_collection is None: -@@ -755,6 +756,9 @@ class _ParamAndGradBuffer: - +@@ -629,6 +631,9 @@ class _ParamAndGradBuffer: + self.data_parallel_world_size = self.data_parallel_group.size() + self.gradient_scaling_factor = gradient_scaling_factor + self.nccl_ub = nccl_ub ++ disable_param_buffers_cpu_backup = ( ++ disable_param_buffers_cpu_backup and self.ddp_config.use_distributed_optimizer ++ ) + + # Data structures to store underlying buckets and relevant indexing data. + self.buckets = [] +@@ -755,6 +760,12 @@ class _ParamAndGradBuffer: + if self.nccl_ub: # If nccl_ub is True, use nccl_allocator to allocate memory for param_data/grad_data. + assert not disable_grad_buffers_cpu_backup, ( + "disable_grad_buffers_cpu_backup is not supported with nccl_ub=True" ++ ) ++ assert not disable_param_buffers_cpu_backup, ( ++ "disable_param_buffers_cpu_backup is not supported with nccl_ub=True" + ) nccl_allocator.init() pool = nccl_allocator.create_nccl_mem_pool( symmetric=not self.ddp_config.disable_symmetric_registration -@@ -773,8 +777,30 @@ class _ParamAndGradBuffer: +@@ -773,19 +784,48 @@ class _ParamAndGradBuffer: torch.distributed.barrier() else: # If nccl_ub is False, mem_alloc_context is nullcontext. + # Individual param/grad contexts below handle TMS regions separately. mem_alloc_context = nullcontext - -+ def _make_no_backup_context(tag, disable): + ++ def _make_no_backup_context(tag, disable, flag_name="disable_grad_buffers_cpu_backup"): + if disable: + try: + from torch_memory_saver import torch_memory_saver + except ImportError as e: + raise ImportError( -+ "disable_grad_buffers_cpu_backup=True requires torch_memory_saver. " ++ f"{flag_name}=True requires torch_memory_saver. " + "Install with: pip install torch-memory-saver" + ) from e -+ + return partial( + torch_memory_saver.region, + tag=tag, + enable_cpu_backup=False, + ) + return nullcontext -+ + grad_mem_alloc_context = _make_no_backup_context( + "grad_buffer", disable_grad_buffers_cpu_backup + ) ++ param_mem_alloc_context = _make_no_backup_context( ++ "param_buffer", disable_param_buffers_cpu_backup, "disable_param_buffers_cpu_backup" ++ ) + with mem_alloc_context(): # For MXFP8 param: Create a shared buffer for param AG and grad RS for memory efficiency # The buffer is mapped to weight gradients whose dtype is either bf16 or FP32. -@@ -803,12 +829,13 @@ class _ParamAndGradBuffer: - device=torch.cuda.current_device(), - requires_grad=False, - ) -- self.grad_data = torch.zeros( + # It can be temporarily reused by param AG. + if self.ddp_config.use_distributed_optimizer and any(is_mxfp8tensor(p) for p in params): +- self.shared_buffer = torch.zeros( - self.numel, - dtype=self.grad_dtype, - device=torch.cuda.current_device(), - requires_grad=False, -- ) -+ with grad_mem_alloc_context(): -+ self.grad_data = torch.zeros( ++ shared_mem_alloc_context = ( ++ param_mem_alloc_context ++ if disable_param_buffers_cpu_backup ++ else grad_mem_alloc_context + ) ++ with shared_mem_alloc_context(): ++ self.shared_buffer = torch.zeros( + self.numel, + dtype=self.grad_dtype, + device=torch.cuda.current_device(), + requires_grad=False, + ) - + # For FP32 weight grads, only half of the buffer is used to store params in bf16. + if self.grad_dtype == torch.float32: + self.param_data = self.shared_buffer[: math.ceil(self.numel / 2)].view( +@@ -797,18 +837,20 @@ class _ParamAndGradBuffer: + else: + # Only re-map param tensors if using distributed optimizer. + if self.ddp_config.use_distributed_optimizer: +- self.param_data = torch.zeros( ++ with param_mem_alloc_context(): ++ self.param_data = torch.zeros( ++ self.numel, ++ dtype=self.param_dtype, ++ device=torch.cuda.current_device(), ++ requires_grad=False, ++ ) ++ with grad_mem_alloc_context(): ++ self.grad_data = torch.zeros( + self.numel, +- dtype=self.param_dtype, ++ dtype=self.grad_dtype, + device=torch.cuda.current_device(), + requires_grad=False, + ) +- self.grad_data = torch.zeros( +- self.numel, +- dtype=self.grad_dtype, +- device=torch.cuda.current_device(), +- requires_grad=False, +- ) + self.grad_data_size = 0 self.param_data_size = 0 diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py @@ -148,15 +196,15 @@ index ef8527e9e..57fbe5bd7 100644 +++ b/megatron/core/extensions/transformer_engine.py @@ -639,6 +639,7 @@ class TELinear(te.pytorch.Linear): self.te_quant_params: Optional[TEQuantizationParams] = None - + for param in self.parameters(): + setattr(param, "parallel_mode", parallel_mode) if is_expert: # Reduce the gradient on the expert_data_parallel group for expert linear layers setattr(param, "allreduce", not self.expert_parallel) @@ -1455,6 +1456,61 @@ class TEDotProductAttention(te.pytorch.DotProductAttention): - - + + if HAVE_TE and is_te_min_version("1.9.0.dev0"): + def ceil_div(x: int, y: int) -> int: + return (x + y - 1) // y @@ -213,13 +261,13 @@ index ef8527e9e..57fbe5bd7 100644 + x_out.main_grad = x.main_grad + + return x_out - + class TEGroupedLinear(te.pytorch.GroupedLinear): """ @@ -1671,6 +1727,20 @@ if HAVE_TE and is_te_min_version("1.9.0.dev0"): return out return out, None - + + def _get_weight_tensors(self): + """Get the weight tensors of the module.""" + weight_tensors = super()._get_weight_tensors() @@ -252,7 +300,7 @@ index 1fd5dcfae..75e1072d5 100644 @@ -434,21 +435,27 @@ def rotary_fwd_kv_kernel( cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - + - KV_ptr = KV + pid_m * stride_kv_seq + pid_head * BLOCK_H * stride_kv_nheads - kv_off = tl.arange(0, BLOCK_H)[:, None] * stride_kv_nheads - mask = kv_off < head_num * stride_kv_nheads @@ -272,12 +320,12 @@ index 1fd5dcfae..75e1072d5 100644 + else: + v = tl.zeros((BLOCK_H, 1), dtype=KV.dtype.element_ty) + k = tl.load(KV_ptr + k_off, mask=mask_k) - + - K_ptr = O_KEY + pid_m * stride_k_seq + pid_head * BLOCK_H * stride_k_nheads - V_ptr = O_VALUE + pid_m * stride_v_seq + pid_head * BLOCK_H * stride_v_nheads + K_ptr = O_KEY + pid_m * stride_k_seq # + pid_head * BLOCK_H * stride_k_nheads + V_ptr = O_VALUE + pid_m * stride_v_seq # + pid_head * BLOCK_H * stride_v_nheads - + - k_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + tl.arange(0, k_dim)[None, :] - v_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_v_nheads + tl.arange(0, v_dim)[None, :] - tl.store(K_ptr + k_out_off, k, mask=mask) @@ -287,13 +335,13 @@ index 1fd5dcfae..75e1072d5 100644 + if v_dim > 0: + v_out_off = ki_range * stride_v_nheads + tl.arange(0, v_dim)[None, :] + tl.store(V_ptr + v_out_off, v, mask=mask_v) - + EMB = K_POS_EMB + pid_m * stride_emb_seq # x1 = t[..., 0::2], x2 = t[..., 1::2] @@ -460,14 +467,16 @@ def rotary_fwd_kv_kernel( x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - + + x_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H + mask_x = x_range < head_num x_left_off = ( @@ -307,8 +355,8 @@ index 1fd5dcfae..75e1072d5 100644 - tl.store(K_ptr + x_right_off, x_right, mask=mask) + tl.store(K_ptr + x_left_off, x_left, mask=mask_x) + tl.store(K_ptr + x_right_off, x_right, mask=mask_x) - - + + @triton.autotune( @@ -493,6 +502,7 @@ def rotary_bwd_kv_kernel( SIN, @@ -321,7 +369,7 @@ index 1fd5dcfae..75e1072d5 100644 @@ -533,27 +543,32 @@ def rotary_bwd_kv_kernel( else: token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) - + - dKV_ptr = dKV + pid_m * stride_dkv_seq + pid_head * BLOCK_H * stride_dkv_nheads - dkv_off = tl.arange(0, BLOCK_H)[:, None] * stride_dkv_nheads - mask = dkv_off < head_num * stride_dkv_nheads @@ -355,7 +403,7 @@ index 1fd5dcfae..75e1072d5 100644 + dv_in_off = ki_range * stride_dv_nheads + tl.arange(0, v_dim)[None, :] + dv = tl.load(dV_ptr + dv_in_off, mask=mask_v) + tl.store(dKV_ptr + dv_out_off, dv, mask=mask_v) - + if pid_head == 0: x_left_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) x_right_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) @@ -368,11 +416,11 @@ index 1fd5dcfae..75e1072d5 100644 x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] x_right_off = x_left_off + emb_dim // 2 @@ -632,6 +647,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function): - + o_key = kv.new_empty(total_seqlen, nheads, emb_dim + k_dim) o_value = kv.new_empty(total_seqlen, nheads, v_dim) + k_dim_ceil = triton.next_power_of_2(k_dim) - + grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) rotary_fwd_kv_kernel[grid]( @@ -643,6 +659,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function): @@ -384,11 +432,11 @@ index 1fd5dcfae..75e1072d5 100644 nheads, batch_size, @@ -700,6 +717,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function): - + d_kv = dk.new_empty(total_seqlen, nheads, ctx.k_dim + ctx.v_dim) d_emb = dk.new_empty(total_seqlen, 1, ctx.emb_dim) + k_dim_ceil = triton.next_power_of_2(ctx.k_dim) - + grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) rotary_bwd_kv_kernel[grid]( @@ -711,6 +729,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function): @@ -406,14 +454,14 @@ index 5dc2d5030..2b241c86d 100644 @@ -61,8 +61,8 @@ except ImportError: try: from torch_memory_saver import torch_memory_saver - + - torch_memory_saver.hook_mode = "torch" - HAVE_TORCH_MEMORY_SAVER = True + # torch_memory_saver.hook_mode = "torch" + HAVE_TORCH_MEMORY_SAVER = False except ImportError: HAVE_TORCH_MEMORY_SAVER = False - + diff --git a/megatron/core/models/common/embeddings/rotary_pos_embedding.py b/megatron/core/models/common/embeddings/rotary_pos_embedding.py index 05a7e8f60..881cfbcaa 100644 --- a/megatron/core/models/common/embeddings/rotary_pos_embedding.py @@ -449,7 +497,7 @@ index 5bb479ad3..a9d3583e5 100755 + post_mlp_layernorm: bool = False, ) -> ModuleSpec: """Use this spec to use lower-level Transformer Engine modules (required for fp8 training). - + @@ -263,9 +265,11 @@ def get_gpt_layer_with_transformer_engine_spec( ), ), @@ -519,7 +567,7 @@ index 5b31ddedf..ead60f2dd 100644 inference_context=inference_context, + mtp_kwargs=mtp_kwargs, ) - + def _postprocess( @@ -581,6 +585,7 @@ class GPTModel(LanguageModule): runtime_gather_output=None, @@ -528,10 +576,10 @@ index 5b31ddedf..ead60f2dd 100644 + mtp_kwargs=None, ): """Postprocesses decoder hidden states to generate logits or compute loss. - + @@ -592,10 +597,12 @@ class GPTModel(LanguageModule): assert runtime_gather_output, "Inference must always gather TP logits" - + # logits and loss + mtp_kwargs = mtp_kwargs or {} + mtp_labels = mtp_kwargs.get('mtp_labels') @@ -546,7 +594,7 @@ index 5b31ddedf..ead60f2dd 100644 @@ -614,13 +621,35 @@ class GPTModel(LanguageModule): if not self.post_process: return hidden_states - + - if self.config.mtp_num_layers is not None: - mtp_labels = labels.clone() + if self.config.mtp_num_layers and mtp_labels is not None: @@ -600,30 +648,58 @@ index a4364f5e9..c76f6daac 100644 param_group["step"] = int(step) + if "step" in param_group and param_group["step"] is None: + del param_group["step"] - + # Grad scaler state. if self.grad_scaler: -@@ -1667,6 +1669,8 @@ class DistributedOptimizer(MixedPrecisionOptimizer): +@@ -969,7 +971,12 @@ class DistributedOptimizer(MixedPrecisionOptimizer): + for bucket_idx, gbuf_range_map in enumerate(gbuf_range_map_for_all_buckets): + bucket_state = [] + for model_param, param_range_map in gbuf_range_map["param_map"].items(): + tensors = self._get_main_param_and_optimizer_states(model_param) ++ if "step" in tensors: ++ # Step is restored from optimizer param_groups. Keeping it in ++ # bucket state makes it common checkpoint state whose list ++ # skeleton depends on save-time optimizer placement. ++ del tensors["step"] + tensors.update( + { + "gbuf_local_start": param_range_map["gbuf_local"].start, +@@ -1667,6 +1669,11 @@ class DistributedOptimizer(MixedPrecisionOptimizer): if key == 'padding': tensors[key] = LocalNonpersistentObject(tensors[key]) continue + if key == 'step': ++ # The optimizer state of STEP is a 0-dim tensor and is handled ++ # separately via param_groups, not as part of the gradient buffer. ++ tensors[key] = LocalNonpersistentObject(tensors[key]) + continue assert tensors[key].shape == (gbuf_local_end - gbuf_local_start,), ( tensors[key].shape, gbuf_local_start, +@@ -1808,6 +1815,11 @@ class DistributedOptimizer(MixedPrecisionOptimizer): + for src_tensors, (model_param, param_range_map) in zip( + bucket_state, gbuf_range_map["param_map"].items() + ): ++ # Local metadata used for checkpoint merging/filtering, not optimizer state. ++ src_tensors.pop('padding', None) ++ # Step is restored from optimizer param_groups. ++ src_tensors.pop('step', None) ++ + # Main param & optimizer states. + self._set_main_param_and_optimizer_states(model_param, src_tensors) + diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py index 7bb964078..2fe9a8cdc 100644 --- a/megatron/core/parallel_state.py +++ b/megatron/core/parallel_state.py @@ -11,6 +11,7 @@ from typing import Callable, List, Optional - + import numpy as np import torch +import torch.distributed as dist - + from .utils import GlobalMemoryBuffer, GlobalSymmetricMemoryBuffer, is_torch_min_version - + diff --git a/megatron/core/pipeline_parallel/p2p_communication.py b/megatron/core/pipeline_parallel/p2p_communication.py index ac839c21f..f18309217 100644 --- a/megatron/core/pipeline_parallel/p2p_communication.py @@ -662,7 +738,7 @@ index 75825cd37..445b3fb84 100644 @@ -711,6 +711,9 @@ def topk_routing_with_score_function( scores, topk, num_groups, group_topk, _compute_topk ) - + + from vime.utils.routing_replay import get_routing_replay_compute_topk + compute_topk = get_routing_replay_compute_topk(compute_topk) + @@ -676,7 +752,7 @@ index a2f3e90bd..b6f732561 100644 @@ -207,6 +207,9 @@ class TopKRouter(Router): if self.config.moe_enable_routing_replay: self.router_replay = RouterReplay() - + + from vime.utils.routing_replay import register_routing_replay + register_routing_replay(self) + @@ -709,17 +785,17 @@ index b0476155a..63f81465d 100755 # embedding decoder_input = embedding(input_ids=input_ids, position_ids=position_ids) + decoder_input = decoder_input.detach() - + - hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + hidden_states = make_viewless_tensor( + inp=hidden_states, requires_grad=True, keep_graph=False + ) - + return input_ids, position_ids, decoder_input, hidden_states - + @@ -821,22 +825,60 @@ class MultiTokenPredictionLayer(MegatronModule): return hidden_states - + def _checkpointed_forward(self, forward_func, *args, **kwargs): + """Wrap forward_func with activation checkpointing while only passing tensors.""" + @@ -764,7 +840,7 @@ index b0476155a..63f81465d 100755 """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`""" if self.config.fp8: from megatron.core.extensions.transformer_engine import te_checkpoint - + return te_checkpoint( - forward_func, + run, @@ -780,7 +856,7 @@ index b0476155a..63f81465d 100755 - forward_func, self.config.distribute_saved_activations, *args, *kwargs.values() + run, self.config.distribute_saved_activations, *tensor_args_tuple ) - + if self.config.recompute_method == 'uniform': diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index dce438520..de51edaf3 100644 @@ -789,13 +865,13 @@ index dce438520..de51edaf3 100644 @@ -229,6 +229,9 @@ class TransformerConfig(ModelParallelConfig): attention_output_gate: bool = False """Whether to apply output gate to the attention layers.""" - + + post_self_attn_layernorm: bool = False + post_mlp_layernorm: bool = False + test_mode: bool = False """Whether to run real-time tests.""" - + diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 12c248684..227e95862 100644 --- a/megatron/core/transformer/transformer_layer.py @@ -805,7 +881,7 @@ index 12c248684..227e95862 100644 self_attention: Union[ModuleSpec, type] = IdentityOp self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp + post_self_attn_layernorm: Union[ModuleSpec, type] = IdentityOp - + pre_cross_attn_layernorm: Union[ModuleSpec, type] = IdentityOp cross_attention: Union[ModuleSpec, type] = IdentityOp @@ -232,6 +233,7 @@ class TransformerLayerSubmodules: @@ -813,13 +889,13 @@ index 12c248684..227e95862 100644 mlp: Union[ModuleSpec, type] = IdentityOp mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp + post_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp - + # Mapping for sharded tensor keys to be applied in `sharded_state_dict` method sharded_state_dict_keys_map: Dict[str, str] = field(default_factory=dict) @@ -311,6 +313,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): # [Module 3: BiasDropoutFusion] self.self_attn_bda = build_module(submodules.self_attn_bda) - + + self.post_self_attn_layernorm = build_module( + submodules.post_self_attn_layernorm, + config=self.config, @@ -831,9 +907,9 @@ index 12c248684..227e95862 100644 self.pre_cross_attn_layernorm = build_module( submodules.pre_cross_attn_layernorm, @@ -376,6 +385,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): - + self.is_moe_layer = isinstance(self.mlp, MoELayer) - + + self.post_mlp_layernorm = build_module( + submodules.post_mlp_layernorm, + config=self.config, @@ -847,7 +923,7 @@ index 12c248684..227e95862 100644 @@ -615,6 +631,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): attention_output_with_bias[0] ) - + + attention_output, attention_output_bias = attention_output_with_bias + attention_output = self.post_self_attn_layernorm(attention_output) + attention_output_with_bias = (attention_output, attention_output_bias) @@ -858,7 +934,7 @@ index 12c248684..227e95862 100644 @@ -794,6 +814,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): self.config.inference_fuse_tp_communication ) - + + mlp_output, mlp_output_bias = mlp_output_with_bias + mlp_output = self.post_mlp_layernorm(mlp_output) + mlp_output_with_bias = (mlp_output, mlp_output_bias) @@ -871,9 +947,9 @@ index 1af066a82..8c7acadb6 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1388,6 +1388,9 @@ def core_transformer_config_from_args(args, config_class=None): - + kw_args['inference_sampling_seed'] = args.seed - + + kw_args['post_self_attn_layernorm'] = args.post_self_attn_layernorm + kw_args['post_mlp_layernorm'] = args.post_mlp_layernorm + @@ -903,10 +979,10 @@ index 17df57dda..260a5f6c8 100644 ) self._vocab = self._tokenizer.get_vocab() diff --git a/megatron/training/training.py b/megatron/training/training.py -index e9736ac08..a83482bf1 100644 +index e9736ac08..6567ed426 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py -@@ -1327,6 +1327,12 @@ def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap +@@ -1327,6 +1327,14 @@ def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap # Wait for the default stream to complete before starting ddp_stream ddp_stream.wait_stream(torch.cuda.current_stream()) # Make ddp_stream start after whatever the default stream already queued @@ -915,11 +991,13 @@ index e9736ac08..a83482bf1 100644 + dp_extra_kwargs['disable_grad_buffers_cpu_backup'] = getattr( + args, 'disable_grad_buffers_cpu_backup', False + ) -+ ++ dp_extra_kwargs['disable_param_buffers_cpu_backup'] = getattr( ++ args, 'disable_param_buffers_cpu_backup', False ++ ) with torch.cuda.stream(ddp_stream): model = [ DP( -@@ -1337,6 +1343,7 @@ def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap +@@ -1337,6 +1345,7 @@ def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap # model chunks is overlapped with compute anyway. disable_bucketing=(model_chunk_idx > 0) or args.overlap_param_gather_with_optimizer_step, diff --git a/docs/zh/developer_guide/install_flashqla.md b/docs/zh/developer_guide/install_flashqla.md new file mode 100644 index 000000000..647f74c53 --- /dev/null +++ b/docs/zh/developer_guide/install_flashqla.md @@ -0,0 +1,28 @@ +# 安装 FlashQLA + +FlashQLA 是 Qwen GDN kernel 的可选运行后端。安装 FlashQLA 后，仍需要在训练命令中显式加入： + +```bash +--qwen-gdn-backend flashqla +``` + +如果不传该参数，Qwen GDN 仍使用默认的 FLA 后端。 + +## 环境要求 + +使用 `--qwen-gdn-backend flashqla` 前，请确认训练节点满足： + +- PyTorch 2.8 或更新版本。 +- CUDA 12.8 或更新版本。 +- NVIDIA SM90 或更新架构 GPU。 +- 所有训练节点都安装了同一套 FlashQLA Python 包。 + +## Docker 镜像 + +标准 CUDA Docker 镜像会默认安装 FlashQLA： + +```bash +docker build \ + -f docker/Dockerfile \ + -t vime:flashqla . +``` diff --git a/tests/test_qwen3_linear_attention_cu_seqlens.py b/tests/test_qwen3_linear_attention_cu_seqlens.py index 1881dfbc5..6abb3cc9d 100644 --- a/tests/test_qwen3_linear_attention_cu_seqlens.py +++ b/tests/test_qwen3_linear_attention_cu_seqlens.py @@ -119,20 +119,39 @@ def load_module(module_name: str): @pytest.mark.unit @pytest.mark.parametrize( - ("module_name", "class_name"), + ("module_name", "class_name", "args", "expected_backend"), [ - ("vime_plugins.models.qwen3_5", "Qwen3_5GatedDeltaNet"), - ("vime_plugins.models.qwen3_next", "Qwen3NextGatedDeltaNet"), + ("vime_plugins.models.qwen3_5", "Qwen3_5GatedDeltaNet", None, "fla"), + ( + "vime_plugins.models.qwen3_5", + "Qwen3_5GatedDeltaNet", + SimpleNamespace(qwen_gdn_backend="flashqla"), + "flashqla", + ), + ("vime_plugins.models.qwen3_next", "Qwen3NextGatedDeltaNet", None, "fla"), + ( + "vime_plugins.models.qwen3_next", + "Qwen3NextGatedDeltaNet", + SimpleNamespace(qwen_gdn_backend="flashqla"), + "flashqla", + ), ], ) -def test_linear_attention_forwards_cu_seqlens_to_chunk_kernel(monkeypatch, module_name: str, class_name: str): +def test_linear_attention_forwards_cu_seqlens_to_chunk_kernel( + monkeypatch, + module_name: str, + class_name: str, + args, + expected_backend: str, +): module = load_module(module_name) monkeypatch.setattr(module.torch.cuda, "current_device", lambda: "cpu") - monkeypatch.setattr(module, "ShortConvolution", FakeShortConvolution) - monkeypatch.setattr(module, "FusedRMSNormGated", FakeFusedRMSNormGated) + monkeypatch.setattr(module, "ShortConvolution", FakeShortConvolution, raising=False) + monkeypatch.setattr(module, "FusedRMSNormGated", FakeFusedRMSNormGated, raising=False) chunk_calls = [] + selected_backends = [] def fake_chunk_gated_delta_rule( q, @@ -152,14 +171,20 @@ def fake_chunk_gated_delta_rule( assert cu_seqlens is not None return torch.zeros_like(v), None - monkeypatch.setattr(module, "chunk_gated_delta_rule", fake_chunk_gated_delta_rule) + def fake_get_chunk_gated_delta_rule(backend): + selected_backends.append(backend) + return fake_chunk_gated_delta_rule - layer = getattr(module, class_name)(make_config(), layer_idx=0) + monkeypatch.setattr(module, "get_chunk_gated_delta_rule", fake_get_chunk_gated_delta_rule) + + layer = getattr(module, class_name)(make_config(), layer_idx=0, args=args) hidden_states = torch.randn(1, 7, 32) cu_seqlens = torch.tensor([0, 3, 7], dtype=torch.int32) output = layer(hidden_states, cu_seqlens=cu_seqlens) + assert selected_backends == [expected_backend] + assert layer.gdn_backend == expected_backend assert output.shape == hidden_states.shape assert len(chunk_calls) == 1 assert torch.equal(chunk_calls[0], cu_seqlens) diff --git a/tests/test_reloadable_process_group_memory_check.py b/tests/test_reloadable_process_group_memory_check.py new file mode 100644 index 000000000..94d9ebb43 --- /dev/null +++ b/tests/test_reloadable_process_group_memory_check.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import pytest + +from vime.utils import reloadable_process_group as rpg + + +@pytest.mark.unit +def test_selected_comm_ops_skip_memory_check(): + skipped_ops = { + "all_gather_into_tensor", + "allgather_into_tensor_coalesced", + "barrier", + "broadcast_object_list", + "reduce_scatter_tensor", + "all_to_all_single", + "isend", + "irecv", + } + checked_ops = { + "all_reduce", + "all_gather", + "broadcast", + "reduce_scatter", + "all_to_all", + "send", + "recv", + "reduce_scatter_tensor_coalesced", + } + + for op_name in skipped_ops: + assert not rpg._should_check_memory_for_comm(op_name) + + for op_name in checked_ops: + assert rpg._should_check_memory_for_comm(op_name) + + +@pytest.mark.unit +def test_wrap_low_level_call_can_skip_available_memory(monkeypatch): + calls = [] + + def fake_available_memory(): + calls.append("available_memory") + return {"free_GB": 100} + + monkeypatch.setattr(rpg, "available_memory", fake_available_memory) + + with rpg._wrap_low_level_call(check_memory=False): + pass + + assert calls == [] + + +@pytest.mark.unit +def test_wrap_low_level_call_checks_available_memory_by_default(monkeypatch): + calls = [] + + def fake_available_memory(): + calls.append("available_memory") + return {"free_GB": 100} + + monkeypatch.setattr(rpg, "available_memory", fake_available_memory) + + with rpg._wrap_low_level_call(): + pass + + assert calls == ["available_memory"] diff --git a/vime/backends/megatron_utils/actor.py b/vime/backends/megatron_utils/actor.py index 165139dce..8b656cb3a 100644 --- a/vime/backends/megatron_utils/actor.py +++ b/vime/backends/megatron_utils/actor.py @@ -97,9 +97,8 @@ def init( self.args, self.model, convert_to_global_name=args.megatron_to_hf_mode == "raw", - translate_gpu_to_cpu=not self.args.enable_weights_backuper, ), - single_tag=None if args.enable_weights_backuper else "actor", + single_tag=None, ) self._active_model_tag: str | None = "actor" self.weights_backuper.backup("actor") @@ -184,6 +183,8 @@ def wake_up(self) -> None: clear_memory() reload_process_groups() + if self.role == "actor": + self._switch_model("actor") print_memory("after wake_up model") def _get_rollout_data(self, rollout_data_ref: Box) -> RolloutBatch: @@ -379,6 +380,7 @@ def train(self, rollout_id: int, rollout_data_ref: Box, external_data=None): result = None if self.args.offload_train: + del rollout_data self.sleep() return result diff --git a/vime/ray/placement_group.py b/vime/ray/placement_group.py index aec0ddc5c..f6bb85891 100644 --- a/vime/ray/placement_group.py +++ b/vime/ray/placement_group.py @@ -1,3 +1,4 @@ +import copy import logging import socket @@ -140,8 +141,11 @@ def create_training_models(args, pgs, rollout_manager): critic_args = ( parse_megatron_role_args(args, args.megatron_config_path, role="critic") if args.megatron_config_path is not None - else args + else copy.deepcopy(args) ) + if args.megatron_config_path is None: + critic_args.disable_param_buffers_cpu_backup = False + critic_model = allocate_train_group( args=critic_args, num_nodes=args.critic_num_nodes, diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index 0bf14b138..453caea70 100644 --- a/vime/utils/arguments.py +++ b/vime/utils/arguments.py @@ -112,6 +112,13 @@ def add_train_arguments(parser): default="thd", help="The qkv layout for Megatron backend.", ) + parser.add_argument( + "--qwen-gdn-backend", + type=str, + choices=["fla", "flashqla"], + default="fla", + help="GDN implementation backend for Qwen linear-attention layers.", + ) parser.add_argument( "--train-env-vars", type=json.loads, @@ -124,12 +131,6 @@ def add_train_arguments(parser): default=1024**3, help="Add margin for train memory allocation. By default we will reserve 1GB as margin.", ) - parser.add_argument( - "--disable-weights-backuper", - action="store_false", - dest="enable_weights_backuper", - help="Whether to disable weights backuper to save host memory.", - ) parser.add_argument( "--megatron-to-hf-mode", choices=["raw", "bridge"], @@ -1502,6 +1503,8 @@ def _apply_megatron_role_overrides(base_args, overrides, role): role_args.use_opd = False role_args.custom_advantage_function_path = None role_args.untie_embeddings_and_output_weights = True + if "disable_param_buffers_cpu_backup" not in overrides: + role_args.disable_param_buffers_cpu_backup = False return role_args @@ -1780,6 +1783,7 @@ def vime_validate_args(args): if args.offload_train: args.disable_grad_buffers_cpu_backup = True + args.disable_param_buffers_cpu_backup = True if args.eval_function_path is None: args.eval_function_path = args.rollout_function_path diff --git a/vime/utils/reloadable_process_group.py b/vime/utils/reloadable_process_group.py index d0ccabfe7..42825e760 100644 --- a/vime/utils/reloadable_process_group.py +++ b/vime/utils/reloadable_process_group.py @@ -11,6 +11,21 @@ old_new_group_dict = {} +_COMM_MEMORY_CHECK_SKIP_OPS = { + "all_gather_into_tensor", + "allgather_into_tensor_coalesced", + "barrier", + "broadcast_object_list", + "reduce_scatter_tensor", + "all_to_all_single", + "isend", + "irecv", +} + + +def _should_check_memory_for_comm(op_name): + return op_name not in _COMM_MEMORY_CHECK_SKIP_OPS + def monkey_patch_torch_dist(): pid = os.getpid() @@ -57,13 +72,14 @@ def new_function(*args, **kwargs): return new_function - def get_new_comm_function(func): + def get_new_comm_function(func, op_name=None): """Wrap communication functions with memory check.""" def new_function(*args, **kwargs): args = tuple([arg.group if isinstance(arg, ReloadableProcessGroup) else arg for arg in args]) kwargs = {k: (v.group if isinstance(v, ReloadableProcessGroup) else v) for k, v in kwargs.items()} - with _wrap_low_level_call(): + check_memory = True if op_name is None else _should_check_memory_for_comm(op_name) + with _wrap_low_level_call(check_memory=check_memory): return func(*args, **kwargs) return new_function @@ -77,19 +93,19 @@ def new_function(*args, **kwargs): dist.all_reduce = get_new_comm_function(dist.all_reduce) dist.all_gather = get_new_comm_function(dist.all_gather) - dist.all_gather_into_tensor = get_new_comm_function(dist.all_gather_into_tensor) + dist.all_gather_into_tensor = get_new_comm_function(dist.all_gather_into_tensor, "all_gather_into_tensor") dist.all_gather_object = get_new_comm_function(dist.all_gather_object) dist.gather_object = get_new_comm_function(dist.gather_object) dist.all_to_all = get_new_comm_function(dist.all_to_all) - dist.all_to_all_single = get_new_comm_function(dist.all_to_all_single) + dist.all_to_all_single = get_new_comm_function(dist.all_to_all_single, "all_to_all_single") dist.broadcast = get_new_comm_function(dist.broadcast) - dist.broadcast_object_list = get_new_comm_function(dist.broadcast_object_list) + dist.broadcast_object_list = get_new_comm_function(dist.broadcast_object_list, "broadcast_object_list") dist.reduce = get_new_comm_function(dist.reduce) dist.reduce_scatter = get_new_comm_function(dist.reduce_scatter) - dist.reduce_scatter_tensor = get_new_comm_function(dist.reduce_scatter_tensor) + dist.reduce_scatter_tensor = get_new_comm_function(dist.reduce_scatter_tensor, "reduce_scatter_tensor") dist.scatter = get_new_comm_function(dist.scatter) dist.gather = get_new_comm_function(dist.gather) - dist.barrier = get_new_comm_function(dist.barrier) + dist.barrier = get_new_comm_function(dist.barrier, "barrier") dist.send = get_new_comm_function(dist.send) dist.recv = get_new_comm_function(dist.recv) dist._coalescing_manager = get_new_comm_function(dist._coalescing_manager) @@ -98,8 +114,8 @@ def new_function(*args, **kwargs): old_isend = dist.isend old_irecv = dist.irecv - dist.isend = get_new_comm_function(dist.isend) - dist.irecv = get_new_comm_function(dist.irecv) + dist.isend = get_new_comm_function(dist.isend, "isend") + dist.irecv = get_new_comm_function(dist.irecv, "irecv") def get_new_p2pop_function(func): def new_function(*args, **kwargs): @@ -192,7 +208,7 @@ def _fwd(self, method, *args, **kwargs): inner = self.group if inner is None: raise RuntimeError("ReloadableProcessGroup: inner PG is None, call reload() first.") - with _wrap_low_level_call(): + with _wrap_low_level_call(check_memory=_should_check_memory_for_comm(method)): return getattr(inner, method)(*args, **kwargs) def _fwd_query(self, method, *args, **kwargs): @@ -294,11 +310,12 @@ def reload_process_groups(): @contextmanager -def _wrap_low_level_call(): +def _wrap_low_level_call(check_memory=True): try: - mem_info = available_memory() - if mem_info["free_GB"] < 3: - clear_memory() + if check_memory: + mem_info = available_memory() + if mem_info["free_GB"] < 3: + clear_memory() yield except Exception as e: mem_info = print_memory("after torch distributed error") diff --git a/vime_plugins/models/qwen3_5.py b/vime_plugins/models/qwen3_5.py index eb8cff0a3..294c9d97a 100644 --- a/vime_plugins/models/qwen3_5.py +++ b/vime_plugins/models/qwen3_5.py @@ -11,11 +11,11 @@ try: from fla.modules import FusedRMSNormGated, ShortConvolution - from fla.ops.gated_delta_rule import chunk_gated_delta_rule except ImportError: pass from .hf_attention import HuggingfaceAttention, _load_hf_config +from .qwen_gdn_backend import get_chunk_gated_delta_rule def _get_text_config(hf_config): @@ -33,8 +33,10 @@ class Qwen3_5GatedDeltaNet(nn.Module): separate in_proj_qkv (for Q,K,V) and in_proj_z (for Z). """ - def __init__(self, config, layer_idx: int): + def __init__(self, config, layer_idx: int, args=None): super().__init__() + self.gdn_backend = getattr(args, "qwen_gdn_backend", "fla") + self.chunk_gated_delta_rule = get_chunk_gated_delta_rule(self.gdn_backend) self.hidden_size = config.hidden_size self.num_v_heads = config.linear_num_value_heads self.num_k_heads = config.linear_num_key_heads @@ -118,7 +120,14 @@ def forward( query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) - core_attn_out, last_recurrent_state = chunk_gated_delta_rule( + if self.gdn_backend == "flashqla": + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + g = g.contiguous() + beta = beta.contiguous() + + core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule( query, key, value, @@ -162,7 +171,7 @@ def __init__( self.hf_config = _get_text_config(self.hf_config) self.hf_config._attn_implementation = "flash_attention_2" - self.linear_attn = Qwen3_5GatedDeltaNet(self.hf_config, self.hf_layer_idx) + self.linear_attn = Qwen3_5GatedDeltaNet(self.hf_config, self.hf_layer_idx, args=args) # Use a simple RMSNorm try: diff --git a/vime_plugins/models/qwen3_next.py b/vime_plugins/models/qwen3_next.py index 3060a945c..683cb2806 100644 --- a/vime_plugins/models/qwen3_next.py +++ b/vime_plugins/models/qwen3_next.py @@ -13,12 +13,12 @@ try: from fla.modules import FusedRMSNormGated, ShortConvolution - from fla.ops.gated_delta_rule import chunk_gated_delta_rule from transformers.models.qwen3_next.modeling_qwen3_next import Qwen3NextAttention, Qwen3NextRMSNorm except ImportError: pass from .hf_attention import HuggingfaceAttention +from .qwen_gdn_backend import get_chunk_gated_delta_rule # adapt from https://github.com/huggingface/transformers/blob/38a08b6e8ae35857109cedad75377997fecbf9d0/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L564 @@ -27,8 +27,10 @@ class Qwen3NextGatedDeltaNet(nn.Module): Qwen3NextGatedDeltaNet with varlen support """ - def __init__(self, config, layer_idx: int): + def __init__(self, config, layer_idx: int, args=None): super().__init__() + self.gdn_backend = getattr(args, "qwen_gdn_backend", "fla") + self.chunk_gated_delta_rule = get_chunk_gated_delta_rule(self.gdn_backend) self.hidden_size = config.hidden_size self.num_v_heads = config.linear_num_value_heads self.num_k_heads = config.linear_num_key_heads @@ -140,7 +142,14 @@ def forward( query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) - core_attn_out, last_recurrent_state = chunk_gated_delta_rule( + if self.gdn_backend == "flashqla": + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + g = g.contiguous() + beta = beta.contiguous() + + core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule( query, key, value, @@ -183,7 +192,7 @@ def __init__( if Qwen3NextAttention is None: raise ImportError("Please install transformers>=4.35.0 to use Qwen3NextAttention.") - self.linear_attn = Qwen3NextGatedDeltaNet(self.hf_config, self.hf_layer_idx) + self.linear_attn = Qwen3NextGatedDeltaNet(self.hf_config, self.hf_layer_idx, args=args) self.input_layernorm = Qwen3NextRMSNorm(self.hf_config.hidden_size, eps=self.hf_config.rms_norm_eps) def hf_forward(self, hidden_states, packed_seq_params): diff --git a/vime_plugins/models/qwen_gdn_backend.py b/vime_plugins/models/qwen_gdn_backend.py new file mode 100644 index 000000000..a95bc8a0d --- /dev/null +++ b/vime_plugins/models/qwen_gdn_backend.py @@ -0,0 +1,46 @@ +import torch + + +def _parse_version(version): + version = version.split("+", 1)[0] + parts = version.split(".") + major = int(parts[0]) + minor = int(parts[1]) if len(parts) > 1 else 0 + return major, minor + + +def _validate_flashqla_runtime(): + if _parse_version(torch.__version__) < (2, 8): + raise RuntimeError(f"FlashQLA backend requires PyTorch 2.8 or newer, got PyTorch {torch.__version__}.") + + if not torch.cuda.is_available(): + raise RuntimeError("FlashQLA backend requires CUDA and an NVIDIA SM90 GPU.") + + major, minor = torch.cuda.get_device_capability() + if (major, minor) < (9, 0): + raise RuntimeError(f"FlashQLA backend requires NVIDIA SM90 or newer, got sm{major}{minor}.") + + cuda_version = torch.version.cuda + if cuda_version is not None and _parse_version(cuda_version) < (12, 8): + raise RuntimeError(f"FlashQLA backend requires CUDA 12.8 or newer, got CUDA {cuda_version}.") + + +def get_chunk_gated_delta_rule(backend: str): + if backend == "fla": + try: + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + except ImportError as exc: + raise ImportError("Qwen GDN backend 'fla' requires flash-linear-attention.") from exc + return chunk_gated_delta_rule + + if backend == "flashqla": + try: + from flash_qla import chunk_gated_delta_rule + except ImportError as exc: + raise ImportError( + "Qwen GDN backend 'flashqla' requires FlashQLA. " "Install it from https://github.com/QwenLM/FlashQLA." + ) from exc + _validate_flashqla_runtime() + return chunk_gated_delta_rule + + raise ValueError(f"Unsupported Qwen GDN backend: {backend}")