|
@@ -0,0 +1,182 @@
|
|
|
|
|
+#!/usr/bin/env python3
|
|
|
|
|
+"""Target-side EAGLE3/aux-hidden-state support for Glm5Next (DFlash2 needs it).
|
|
|
|
|
+
|
|
|
|
|
+The DFlash2 drafter consumes the TARGET's hidden states at layers
|
|
|
|
|
+5/14/24/33/42 (draft config.json dflash_config.target_layer_ids). The vLLM
|
|
|
|
|
+spec-decode runner gates that on supports_eagle3(target) and on the model
|
|
|
|
|
+forward emitting aux states (EagleModelMixin). The day-0 nvidia Glm5Next
|
|
|
|
|
+model implements neither; the fork (local-inference-lab/vllm
|
|
|
|
|
+dev/jovian-judgement @ da4d7be6, Apache-2.0, see UPSTREAM.md) does. A full
|
|
|
|
|
+file swap is NOT possible: the fork diverges from our day-0 base in ~150
|
|
|
|
|
+lines (b12x backends, weight remaps) and our base carries fixes the fork
|
|
|
|
|
+lacks. So this is a surgical port of exactly the fork's:
|
|
|
|
|
+
|
|
|
|
|
+ * EagleModelMixin on Glm5NextModel + SupportsEagle3 on both roots
|
|
|
|
|
+ * dflash_capture flag (fork: speculative_config.use_dflash(), absent on
|
|
|
|
|
+ our base -> method == "dflash" equivalent)
|
|
|
|
|
+ * _prepare_aux_hidden_state (mHC: hc_post then hc_contract for dflash)
|
|
|
|
|
+ * forward-loop aux collection + tuple return (after self.norm, fork parity)
|
|
|
|
|
+
|
|
|
|
|
+Idempotent; every anchor asserted exactly-once; py_compile-checked.
|
|
|
|
|
+"""
|
|
|
|
|
+
|
|
|
|
|
+import py_compile
|
|
|
|
|
+
|
|
|
|
|
+B = "/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia"
|
|
|
|
|
+PATH = f"{B}/model.py"
|
|
|
|
|
+
|
|
|
|
|
+EDITS = [
|
|
|
|
|
+ # 1. imports
|
|
|
|
|
+ (
|
|
|
|
|
+ """from vllm.model_executor.models.interfaces import (
|
|
|
|
|
+ HasInnerState,
|
|
|
|
|
+ IsHybrid,
|
|
|
|
|
+ MixtureOfExperts,
|
|
|
|
|
+ SupportsPP,
|
|
|
|
|
+)""",
|
|
|
|
|
+ """from vllm.model_executor.models.interfaces import (
|
|
|
|
|
+ EagleModelMixin,
|
|
|
|
|
+ HasInnerState,
|
|
|
|
|
+ IsHybrid,
|
|
|
|
|
+ MixtureOfExperts,
|
|
|
|
|
+ SupportsEagle3,
|
|
|
|
|
+ SupportsPP,
|
|
|
|
|
+)""",
|
|
|
|
|
+ ),
|
|
|
|
|
+ # 2. model gets the mixin (runner sets aux layers through it)
|
|
|
|
|
+ (
|
|
|
|
|
+ "class Glm5NextModel(nn.Module):\n",
|
|
|
|
|
+ "class Glm5NextModel(nn.Module, EagleModelMixin):\n",
|
|
|
|
|
+ ),
|
|
|
|
|
+ # 3. dflash capture flag in Glm5NextModel.__init__
|
|
|
|
|
+ (
|
|
|
|
|
+ """ self.is_sequence_parallel = (
|
|
|
|
|
+ vllm_config.parallel_config.use_sequence_parallel_moe
|
|
|
|
|
+ )
|
|
|
|
|
+""",
|
|
|
|
|
+ """ self.is_sequence_parallel = (
|
|
|
|
|
+ vllm_config.parallel_config.use_sequence_parallel_moe
|
|
|
|
|
+ )
|
|
|
|
|
+ _spec = vllm_config.speculative_config
|
|
|
|
|
+ # fork parity: dflash drafts consume hc_contract-ed aux states
|
|
|
|
|
+ # (SpeculativeConfig.use_dflash() only exists on the fork).
|
|
|
|
|
+ self.dflash_capture = (
|
|
|
|
|
+ _spec is not None and _spec.method == "dflash"
|
|
|
|
|
+ )
|
|
|
|
|
+""",
|
|
|
|
|
+ ),
|
|
|
|
|
+ # 4. aux extraction method (fork body, verbatim apart from type-name quote)
|
|
|
|
|
+ (
|
|
|
|
|
+ """ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
|
|
|
|
+ return self.embed_tokens(input_ids)
|
|
|
|
|
+""",
|
|
|
|
|
+ """ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
|
|
|
|
+ return self.embed_tokens(input_ids)
|
|
|
|
|
+
|
|
|
|
|
+ def _prepare_aux_hidden_state(
|
|
|
|
|
+ self,
|
|
|
|
|
+ layer: "Glm5NextDecoderLayer",
|
|
|
|
|
+ hidden_states: torch.Tensor,
|
|
|
|
|
+ residual: torch.Tensor | None,
|
|
|
|
|
+ post: torch.Tensor | None,
|
|
|
|
|
+ comb: torch.Tensor | None,
|
|
|
|
|
+ ) -> torch.Tensor:
|
|
|
|
|
+ if not layer.mhc or residual is None:
|
|
|
|
|
+ return hidden_states
|
|
|
|
|
+ assert post is not None and comb is not None
|
|
|
|
|
+ aux_hidden_state = layer.hc_post(hidden_states, residual, post, comb)
|
|
|
|
|
+ if self.dflash_capture:
|
|
|
|
|
+ return hc_contract(aux_hidden_state, layer.n)
|
|
|
|
|
+ return aux_hidden_state.flatten(1)
|
|
|
|
|
+""",
|
|
|
|
|
+ ),
|
|
|
|
|
+ # 5. forward return type
|
|
|
|
|
+ (
|
|
|
|
|
+ """ ) -> torch.Tensor:
|
|
|
|
|
+ if get_pp_group().is_first_rank:
|
|
|
|
|
+""",
|
|
|
|
|
+ """ ) -> (
|
|
|
|
|
+ torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]
|
|
|
|
|
+ ):
|
|
|
|
|
+ if get_pp_group().is_first_rank:
|
|
|
|
|
+""",
|
|
|
|
|
+ ),
|
|
|
|
|
+ # 6. aux collection in the layer loop (fork semantics: capture the state
|
|
|
|
|
+ # ENTERING layer_idx+1, i.e. the OUTPUT of layer_idx)
|
|
|
|
|
+ (
|
|
|
|
|
+ """ for layer in self._active_layers:
|
|
|
|
|
+ hidden_states, residual, post, comb = layer(
|
|
|
|
|
+ positions, hidden_states, residual, post, comb
|
|
|
|
|
+ )
|
|
|
|
|
+""",
|
|
|
|
|
+ """ aux_hidden_states: list[torch.Tensor] = []
|
|
|
|
|
+ if self.start_layer in self.aux_hidden_state_layers:
|
|
|
|
|
+ aux_hidden_states.append(hidden_states)
|
|
|
|
|
+
|
|
|
|
|
+ for layer_idx, layer in enumerate(
|
|
|
|
|
+ self._active_layers, start=self.start_layer
|
|
|
|
|
+ ):
|
|
|
|
|
+ if layer_idx + 1 in self.aux_hidden_state_layers:
|
|
|
|
|
+ aux_hidden_state = self._prepare_aux_hidden_state(
|
|
|
|
|
+ layer, hidden_states, residual, post, comb
|
|
|
|
|
+ )
|
|
|
|
|
+ if self.is_sequence_parallel:
|
|
|
|
|
+ aux_hidden_state = sp_all_gather(aux_hidden_state)[
|
|
|
|
|
+ :full_num_tokens
|
|
|
|
|
+ ]
|
|
|
|
|
+ aux_hidden_states.append(aux_hidden_state)
|
|
|
|
|
+ hidden_states, residual, post, comb = layer(
|
|
|
|
|
+ positions, hidden_states, residual, post, comb
|
|
|
|
|
+ )
|
|
|
|
|
+""",
|
|
|
|
|
+ ),
|
|
|
|
|
+ # 7. tuple return after norm (fork parity — logits stay well-formed)
|
|
|
|
|
+ (
|
|
|
|
|
+ """ hidden_states = self.norm(hidden_states)
|
|
|
|
|
+ return hidden_states
|
|
|
|
|
+
|
|
|
|
|
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
|
|
|
+ stacked_params_mapping = [
|
|
|
|
|
+""",
|
|
|
|
|
+ """ hidden_states = self.norm(hidden_states)
|
|
|
|
|
+ if aux_hidden_states:
|
|
|
|
|
+ return hidden_states, aux_hidden_states
|
|
|
|
|
+ return hidden_states
|
|
|
|
|
+
|
|
|
|
|
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
|
|
|
+ stacked_params_mapping = [
|
|
|
|
|
+""",
|
|
|
|
|
+ ),
|
|
|
|
|
+ # 8. both registered roots advertise eagle3 (the runner checks the root)
|
|
|
|
|
+ (
|
|
|
|
|
+ """ nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid
|
|
|
|
|
+):
|
|
|
|
|
+""",
|
|
|
|
|
+ """ nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid,
|
|
|
|
|
+ SupportsEagle3
|
|
|
|
|
+):
|
|
|
|
|
+""",
|
|
|
|
|
+ ),
|
|
|
|
|
+ (
|
|
|
|
|
+ """ Glm4vForConditionalGeneration, HasInnerState, IsHybrid
|
|
|
|
|
+):
|
|
|
|
|
+""",
|
|
|
|
|
+ """ Glm4vForConditionalGeneration, HasInnerState, IsHybrid, SupportsEagle3
|
|
|
|
|
+):
|
|
|
|
|
+""",
|
|
|
|
|
+ ),
|
|
|
|
|
+]
|
|
|
|
|
+
|
|
|
|
|
+src = open(PATH).read()
|
|
|
|
|
+if "dflash_capture" in src:
|
|
|
|
|
+ print(f"already patched: {PATH}")
|
|
|
|
|
+else:
|
|
|
|
|
+ for i, (old, new) in enumerate(EDITS, 1):
|
|
|
|
|
+ n = src.count(old)
|
|
|
|
|
+ assert n == 1, f"edit {i}: anchor {n}x in {PATH} — review me"
|
|
|
|
|
+ src = src.replace(old, new, 1)
|
|
|
|
|
+ open(PATH, "w").write(src)
|
|
|
|
|
+ py_compile.compile(PATH, doraise=True)
|
|
|
|
|
+ print(f"patched: {PATH}")
|
|
|
|
|
+
|
|
|
|
|
+print("patch_eagle3 OK")
|