diff --git a/init/dflash100k_epoch6/model.safetensors b/init/dflash100k_epoch6/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..192b09729dafb7a009014e667e4079db43b3eaee --- /dev/null +++ b/init/dflash100k_epoch6/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1fc0023cdb9cd6b93336abbb5a36f8f08527368002d68ca9a1a32b72aeb93073 +size 1074860568 diff --git a/runs/dspark_a/export/epoch_1/head.safetensors b/runs/dspark_a/export/epoch_1/head.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..29510b5bbdd991b58d8c3864b5261a83aebd403b --- /dev/null +++ b/runs/dspark_a/export/epoch_1/head.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:af98cf412711c7d9f7df177dd1bb20b6576bb8233b9d14da5b0b6d3aea24a717 +size 311165128 diff --git a/runs/dspark_joint_6p/ckpt/epoch_6_step_17156.pt b/runs/dspark_joint_6p/ckpt/epoch_6_step_17156.pt new file mode 100644 index 0000000000000000000000000000000000000000..60d7733f2a6b1de1a77afdeb7a067a87ddc02904 --- /dev/null +++ b/runs/dspark_joint_6p/ckpt/epoch_6_step_17156.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4c29ef39cad74a7e958f4c18a31d77e97080db7c1033be6012732e84400b1fd8 +size 7382703887 diff --git a/runs/dspark_joint_6p/export/epoch_6/head.safetensors b/runs/dspark_joint_6p/export/epoch_6/head.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0ea5b4738a62eb0ac6dcc7ef6025453ad8a4a1d0 --- /dev/null +++ b/runs/dspark_joint_6p/export/epoch_6/head.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e740457069c8cdaf7feebc6f1c9a45af543c8208fa4de0aea29af7a8cf70ddaa +size 311165128 diff --git a/runs/dspark_joint_6p/export/epoch_6/model.safetensors b/runs/dspark_joint_6p/export/epoch_6/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..bb2befdde769a6d65842096e5884fa15e93e1ed5 --- /dev/null +++ b/runs/dspark_joint_6p/export/epoch_6/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fcb5d93581ed4fe4cd7405bcebdbf0fe358b8bf0ff2248d2b9e681eba602ddc9 +size 1074860568 diff --git a/runs/dspark_joint_6p_tv/export/epoch_1/config.json b/runs/dspark_joint_6p_tv/export/epoch_1/config.json new file mode 100644 index 0000000000000000000000000000000000000000..1d3773589fd48604b18e25ac315aabd4baacee8e --- /dev/null +++ b/runs/dspark_joint_6p_tv/export/epoch_1/config.json @@ -0,0 +1,55 @@ +{ + "architectures": [ + "DFlashDraftModel" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": { + "AutoModel": "dflash.DFlashDraftModel" + }, + "block_size": 16, + "bos_token_id": 151643, + "dflash_config": { + "mask_token_id": 151669, + "target_layer_ids": [ + 1, + 9, + 17, + 25, + 33 + ] + }, + "dtype": "float32", + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 9728, + "layer_types": [ + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention" + ], + "max_position_embeddings": 40960, + "max_window_layers": 5, + "model_type": "qwen3", + "num_attention_heads": 32, + "num_hidden_layers": 5, + "num_key_value_heads": 8, + "num_target_layers": 36, + "pad_token_id": null, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "rope_theta": 1000000, + "rope_type": "default" + }, + "sliding_window": null, + "tie_word_embeddings": true, + "transformers_version": "5.15.0", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936 +} diff --git a/runs/dspark_joint_6p_tv/export/epoch_1/dspark_config.json b/runs/dspark_joint_6p_tv/export/epoch_1/dspark_config.json new file mode 100644 index 0000000000000000000000000000000000000000..f55bca02da0012c35ca0d03d57fbb66c8eb4b9a9 --- /dev/null +++ b/runs/dspark_joint_6p_tv/export/epoch_1/dspark_config.json @@ -0,0 +1,16 @@ +{ + "rank": 256, + "backbone": { + "kind": "dflash", + "path": "runs/dspark_joint_6p_tv/export/epoch_1" + }, + "epoch": 1, + "step": 2859, + "label": "epoch_1", + "joint": true, + "objective": "dflash", + "kd": true, + "tv": 0.9, + "head_rank": 256, + "loss": "DSpark baseline: joint DFlash + Markov head, (1 - tv) CE + tv L1 against Qwen3-4B probabilities, tv = 0.9" +} \ No newline at end of file diff --git a/runs/dspark_joint_6p_tv/export/epoch_1/head.safetensors b/runs/dspark_joint_6p_tv/export/epoch_1/head.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..bf1bbbb2c6312d000426289fb8c25858eeb3a165 --- /dev/null +++ b/runs/dspark_joint_6p_tv/export/epoch_1/head.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ebfa3ff872cf4defd1c861d422cb4565c3a288ba4862f7a1c4044463e7282e3d +size 311165128 diff --git a/runs/dspark_joint_6p_tv/export/epoch_1/model.safetensors b/runs/dspark_joint_6p_tv/export/epoch_1/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..8ac53f7dbc5632316c91383ff882827527b37a49 --- /dev/null +++ b/runs/dspark_joint_6p_tv/export/epoch_1/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a6e2aa3c88581da6c7c3cacb35f50f678b5fb164f047299deb0f2bc0dbe39589 +size 1074860568 diff --git a/runs/dspark_joint_6p_tv/export/epoch_1/training_recipe.json b/runs/dspark_joint_6p_tv/export/epoch_1/training_recipe.json new file mode 100644 index 0000000000000000000000000000000000000000..a63f1738910da42f215256278c8746fef28da09d --- /dev/null +++ b/runs/dspark_joint_6p_tv/export/epoch_1/training_recipe.json @@ -0,0 +1,10 @@ +{ + "loss": "DSpark baseline: joint DFlash + Markov head, (1 - tv) CE + tv L1 against Qwen3-4B probabilities, tv = 0.9", + "epoch": 1, + "step": 2859, + "label": "epoch_1", + "objective": "dflash", + "kd": true, + "tv": 0.9, + "head_rank": 256 +} \ No newline at end of file diff --git a/runs/fmlm2_v3/export/epoch_1/flow.safetensors b/runs/fmlm2_v3/export/epoch_1/flow.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..6f8a19cafae852cfd57b46c61d97ac85fd960322 --- /dev/null +++ b/runs/fmlm2_v3/export/epoch_1/flow.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:38594ad798d0869e11e0b7f1aadfeb8fb075f0739a489e129cd28e8ffac4f5f1 +size 1770876772 diff --git a/runs/fmlm2_v3/export/epoch_1/head.safetensors b/runs/fmlm2_v3/export/epoch_1/head.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..9b9254b04fa46c6698cf389941997e89a4140468 --- /dev/null +++ b/runs/fmlm2_v3/export/epoch_1/head.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:acb621a7e9913a6e9df3e151222a2b0110b2b06cb7b9f6e37c225545b401bd55 +size 311165128 diff --git a/runs/fmlm2_v3/export/epoch_1/model.safetensors b/runs/fmlm2_v3/export/epoch_1/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..96f8f20224a2e124b2d0be293e0273052d087746 --- /dev/null +++ b/runs/fmlm2_v3/export/epoch_1/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:be6d410548cea0f60bdde21e66fc073ae9d2aa059160c8bfad6a64fd366950cf +size 1074860568 diff --git a/runs/fmlm2_v3_6ep/ckpt/epoch_6_step_17156.pt b/runs/fmlm2_v3_6ep/ckpt/epoch_6_step_17156.pt new file mode 100644 index 0000000000000000000000000000000000000000..67fba94638abe3f4f1d5b4694886fd038fa62f62 --- /dev/null +++ b/runs/fmlm2_v3_6ep/ckpt/epoch_6_step_17156.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1aaec2649c7dc202b6931d991adafb7d52a7e354bdef79920cca2a06ad8a66da +size 12695355007 diff --git a/runs/fmlm2_v3_6ep/export/epoch_6/flow.safetensors b/runs/fmlm2_v3_6ep/export/epoch_6/flow.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..df2d9e9001d0407bc3774fe33618a0a53ad2c8cc --- /dev/null +++ b/runs/fmlm2_v3_6ep/export/epoch_6/flow.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:41f71c01d20ad59568b67a80e24b338f31e3ff71f3c9f16bcc60977b823ee340 +size 1770876772 diff --git a/runs/fmlm2_v3_6ep/export/epoch_6/head.safetensors b/runs/fmlm2_v3_6ep/export/epoch_6/head.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..54d2935fd3cfd0061355f8725a46ab834a43ea61 --- /dev/null +++ b/runs/fmlm2_v3_6ep/export/epoch_6/head.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:52f3ad42a9fb35d38661d0f09a7edc8e30220c8a0435f3ad8ebe70915fa683f5 +size 311165128 diff --git a/runs/fmlm2_v3_6ep/export/epoch_6/model.safetensors b/runs/fmlm2_v3_6ep/export/epoch_6/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..fc38c9cf0b4b089bc4ca85d8a8276468b181e2cf --- /dev/null +++ b/runs/fmlm2_v3_6ep/export/epoch_6/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:066d537ccc3da6eb3da7b6855b6d3f226ebd89dbd0280168fab30e3f0b3c3ea7 +size 1074860568 diff --git a/runs/fmlm2_v3_6ep_bv/export/epoch_1/flow.safetensors b/runs/fmlm2_v3_6ep_bv/export/epoch_1/flow.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e9e8b66128bcab4479cb8c97bc56779a16349538 --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv/export/epoch_1/flow.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4fc7410e6b5a259a6006266d871eea6fc7fce72eb9da4f1ebcee53c62eac0aa3 +size 1770876772 diff --git a/runs/fmlm2_v3_6ep_bv/export/epoch_1/head.safetensors b/runs/fmlm2_v3_6ep_bv/export/epoch_1/head.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..63eb6e9018937d9e8ae6c47946b4b3808c63b82d --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv/export/epoch_1/head.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:252f749694020bf01d7ed7f991776f852a25aea578ad72e431f9b8ea2e3fc329 +size 311165128 diff --git a/runs/fmlm2_v3_6ep_bv/export/epoch_1/model.safetensors b/runs/fmlm2_v3_6ep_bv/export/epoch_1/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..fc9dbe875d9a00ad40c412e29e1cf2bd9d767274 --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv/export/epoch_1/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6a3d2040a940e3526e7e9df6abea13efebf51d377c942655545e3e8971249d20 +size 1074860568 diff --git a/runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/flow.safetensors b/runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/flow.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..11971575e46c863a1c2157bfdf499982f80f3199 --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/flow.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b3fd239bb0f2e19eee270407c0c41fe75012f6e9f9270559d2426e019ce8cec9 +size 1770876772 diff --git a/runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/head.safetensors b/runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/head.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..8787a04798f72fc1eb0b5006b2424bda2532203d --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/head.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b10c48c786c1ff5aa26624953174f48feba3ad2c10b31143691a4e91fa071806 +size 311165128 diff --git a/runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/model.safetensors b/runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..81fd02af53b072e288f481bed488ae289c98919a --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fbc731478b05f229b4e00afed51d29e452f45ced4f588f4e8f85413a05eeeded +size 1074860568 diff --git a/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/config.json b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/config.json new file mode 100644 index 0000000000000000000000000000000000000000..89073407b95c2f29693c309b6f00189c99b32b61 --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/config.json @@ -0,0 +1,52 @@ +{ + "architectures": [ + "DFlashDraftModel" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": { + "AutoModel": "dflash.DFlashDraftModel" + }, + "block_size": 16, + "bos_token_id": 151643, + "dflash_config": { + "mask_token_id": 151669, + "target_layer_ids": [ + 1, + 9, + 17, + 25, + 33 + ] + }, + "dtype": "bfloat16", + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 9728, + "layer_types": [ + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention" + ], + "max_position_embeddings": 40960, + "max_window_layers": 5, + "model_type": "qwen3", + "num_attention_heads": 32, + "num_hidden_layers": 5, + "num_key_value_heads": 8, + "num_target_layers": 36, + "rms_norm_eps": 1e-06, + "rope_scaling": null, + "rope_theta": 1000000, + "sliding_window": null, + "tie_word_embeddings": true, + "transformers_version": "4.57.3", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936 +} diff --git a/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/dflash.py b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/dflash.py new file mode 100644 index 0000000000000000000000000000000000000000..f3cd5d26057e92f4721977cf53195c2d738a6b76 --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/dflash.py @@ -0,0 +1,277 @@ +from typing import Optional, Callable +from typing_extensions import Unpack, Tuple +import torch +from torch import nn +from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3RMSNorm, + Qwen3RotaryEmbedding, + Qwen3Config, + Qwen3PreTrainedModel, + Qwen3MLP, + GradientCheckpointingLayer, + FlashAttentionKwargs, + rotate_half, + eager_attention_forward, + ALL_ATTENTION_FUNCTIONS, +) +from transformers import DynamicCache +from transformers.modeling_outputs import CausalLMOutputWithPast +from transformers.cache_utils import Cache +from .utils import build_target_layer_ids, extract_context_feature, sample + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_len = q.size(-2) + q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :]) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + +class Qwen3DFlashAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = False + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None + + def forward( + self, + hidden_states: torch.Tensor, + target_hidden: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_values: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + bsz, q_len = hidden_states.shape[:-1] + ctx_len = target_hidden.shape[1] + q = self.q_proj(hidden_states) + q = q.view(bsz, q_len, -1, self.head_dim) + q = self.q_norm(q).transpose(1, 2) + k_ctx = self.k_proj(target_hidden) + k_noise = self.k_proj(hidden_states) + v_ctx = self.v_proj(target_hidden) + v_noise = self.v_proj(hidden_states) + k = torch.cat([k_ctx, k_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim) + v = torch.cat([v_ctx, v_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim) + k = self.k_norm(k).transpose(1, 2) + v = v.transpose(1, 2) + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb(q, k, cos, sin) + if past_key_values is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs) + attn_fn: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + attn_output, attn_weights = attn_fn( + self, + q, + k, + v, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, q_len, -1) + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + +class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Qwen3DFlashAttention(config=config, layer_idx=layer_idx) + self.mlp = Qwen3MLP(config) + self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + target_hidden: Optional[torch.Tensor] = None, + hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + hidden_states=hidden_states, + target_hidden=target_hidden, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + )[0] + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + +class DFlashDraftModel(Qwen3PreTrainedModel): + config_class = Qwen3Config + _no_split_modules = ["Qwen3DFlashDecoderLayer"] + + def __init__(self, config) -> None: + super().__init__(config) + self.config = config + self.layers = nn.ModuleList( + [Qwen3DFlashDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.target_layer_ids = self.config.dflash_config.get("target_layer_ids", build_target_layer_ids(config.num_target_layers, config.num_hidden_layers)) + self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3RotaryEmbedding(config) + self.fc = nn.Linear(len(self.target_layer_ids) * config.hidden_size, config.hidden_size, bias=False) + self.hidden_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.block_size = config.block_size + self.mask_token_id = self.config.dflash_config.get("mask_token_id", None) + self.post_init() + + def forward( + self, + position_ids: torch.LongTensor, + attention_mask: Optional[torch.Tensor] = None, + noise_embedding: Optional[torch.Tensor] = None, + target_hidden: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: bool = False, + **kwargs, + ) -> CausalLMOutputWithPast: + hidden_states = noise_embedding + target_hidden = self.hidden_norm(self.fc(target_hidden)) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for layer in self.layers: + hidden_states = layer( + hidden_states=hidden_states, + target_hidden=target_hidden, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + return self.norm(hidden_states) + + @torch.inference_mode() + def spec_generate( + self, + target: nn.Module, + input_ids: torch.LongTensor, + max_new_tokens: int, + stop_token_ids: list[int], + temperature: float, + ): + self.eval() + num_input_tokens = input_ids.shape[1] + max_length = num_input_tokens + max_new_tokens + + block_size = self.block_size + output_ids = torch.full( + (1, max_length + block_size), + self.mask_token_id, + dtype=torch.long, + device=target.device, + ) + position_ids = torch.arange(output_ids.shape[1], device=target.device).unsqueeze(0) + + past_key_values_target = DynamicCache() + past_key_values_draft = DynamicCache() + + # Prefill stage + output = target( + input_ids, + position_ids=position_ids[:, :num_input_tokens], + past_key_values=past_key_values_target, + use_cache=True, + logits_to_keep=1, + output_hidden_states=True, + ) + + output_ids[:, :num_input_tokens] = input_ids + output_ids[:, num_input_tokens:num_input_tokens+1] = sample(output.logits, temperature) + target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids) + + # Decode stage + acceptance_lengths = [] + start = input_ids.shape[1] + while start < max_length: + block_output_ids = output_ids[:, start : start + block_size].clone() + block_position_ids = position_ids[:, start : start + block_size] + noise_embedding = target.model.embed_tokens(block_output_ids) + draft_logits = target.lm_head(self( + target_hidden=target_hidden, + noise_embedding=noise_embedding, + position_ids=position_ids[:, past_key_values_draft.get_seq_length(): start + block_size], + past_key_values=past_key_values_draft, + use_cache=True, + is_causal=False, + )[:, -block_size+1:, :]) + past_key_values_draft.crop(start) + block_output_ids[:, 1:] = sample(draft_logits) + + output = target( + block_output_ids, + position_ids=block_position_ids, + past_key_values=past_key_values_target, + use_cache=True, + output_hidden_states=True, + ) + + posterior = sample(output.logits, temperature) + acceptance_length = (block_output_ids[:, 1:] == posterior[:, :-1]).cumprod(dim=1).sum(dim=1)[0].item() + output_ids[:, start : start + acceptance_length + 1] = block_output_ids[:, : acceptance_length + 1] + output_ids[:, start + acceptance_length + 1] = posterior[:, acceptance_length] + start += acceptance_length + 1 + past_key_values_target.crop(start) + target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids)[:, :acceptance_length + 1, :] + acceptance_lengths.append(acceptance_length+1) + if stop_token_ids is not None and any( + stop_token_id in output_ids[:, num_input_tokens:] for stop_token_id in stop_token_ids + ): + break + output_ids = output_ids[:, :max_length] + output_ids = output_ids[:, output_ids[0] != self.mask_token_id] + if stop_token_ids is not None: + stop_token_ids = torch.tensor(stop_token_ids, device=output_ids.device) + stop_token_indices = torch.isin(output_ids[0][num_input_tokens:], stop_token_ids).nonzero(as_tuple=True)[0] + if stop_token_indices.numel() > 0: + output_ids = output_ids[:, : num_input_tokens + stop_token_indices[0] + 1] + + return output_ids diff --git a/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/dspark_config.json b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/dspark_config.json new file mode 100644 index 0000000000000000000000000000000000000000..3f20eb8b777c590634f3db1d28fac4aec1055fe7 --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/dspark_config.json @@ -0,0 +1,21 @@ +{ + "rank": 256, + "backbone": { + "kind": "flow", + "path": "runs/fmlm2_v3_6ep_bv_lt/export/epoch_1", + "mode": "onejump_zero" + }, + "epoch": 1, + "step": 2859, + "label": "epoch_1", + "joint": true, + "objective": "fmlm2", + "kd": true, + "readout_frac": 0.25, + "head_rank": 256, + "tv": 0.0, + "bv": 0.9, + "lafa_time": true, + "emb_mse": 0.0, + "fix_embed_norm": false +} \ No newline at end of file diff --git a/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/flow.safetensors b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/flow.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a0c09e8f66f32775c4df1d0a98cff099eca6f343 --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/flow.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:40b4b27ad2a1a5a4fa1fca89b55b411b39a0beaa4c34d4c25335786ab1256ef9 +size 1770876772 diff --git a/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/flow_config.json b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/flow_config.json new file mode 100644 index 0000000000000000000000000000000000000000..a8fa8900a7983a608256f87200e526c4f75934a5 --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/flow_config.json @@ -0,0 +1,15 @@ +{ + "learn_embed": true, + "epoch": 1, + "step": 2859, + "label": "epoch_1", + "objective": "fmlm2", + "kd": true, + "readout_frac": 0.25, + "head_rank": 256, + "tv": 0.0, + "bv": 0.9, + "lafa_time": true, + "emb_mse": 0.0, + "fix_embed_norm": false +} \ No newline at end of file diff --git a/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/head.safetensors b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/head.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..8d4ccde273ba92cbd70fa9d28cfa7bf75a88ac6c --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/head.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c15930d6cff545a25aad05afe29f81e89b750f1eb4705191473e1d12d8f85e9b +size 311165128 diff --git a/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/model.safetensors b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..1cf67b1297b3c783bafc3d6b7f6ab4f8a9bfbe7c --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c9b9a73b54c60f89efab5ce15c70454367ec1b30cc328bbc25cc5f9563d6fcb5 +size 1074860568 diff --git a/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/modeling_dflash.py b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/modeling_dflash.py new file mode 100644 index 0000000000000000000000000000000000000000..ec9f19ffd11e242b0dd9e9469817f21ef31c2b04 --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/modeling_dflash.py @@ -0,0 +1,309 @@ +from typing import Optional, Callable +from typing_extensions import Unpack, Tuple +import torch +from torch import nn +from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3RMSNorm, + Qwen3RotaryEmbedding, + Qwen3Config, + Qwen3PreTrainedModel, + Qwen3MLP, + GradientCheckpointingLayer, + FlashAttentionKwargs, + rotate_half, + eager_attention_forward, + ALL_ATTENTION_FUNCTIONS, +) +from transformers import DynamicCache +from transformers.modeling_outputs import CausalLMOutputWithPast +from transformers.cache_utils import Cache + +def sample(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor: + if temperature < 1e-5: + return torch.argmax(logits, dim=-1) + bsz, seq_len, vocab_size = logits.shape + logits = logits.view(-1, vocab_size) + logits = logits / temperature + probs = torch.softmax(logits, dim=-1) + return torch.multinomial(probs, num_samples=1).view(bsz, seq_len) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_len = q.size(-2) + q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :]) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + +class Qwen3DFlashAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = False + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None + + def forward( + self, + hidden_states: torch.Tensor, + target_hidden: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_values: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + bsz, q_len = hidden_states.shape[:-1] + ctx_len = target_hidden.shape[1] + q = self.q_proj(hidden_states) + q = q.view(bsz, q_len, -1, self.head_dim) + q = self.q_norm(q).transpose(1, 2) + k_ctx = self.k_proj(target_hidden) + k_noise = self.k_proj(hidden_states) + v_ctx = self.v_proj(target_hidden) + v_noise = self.v_proj(hidden_states) + k = torch.cat([k_ctx, k_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim) + v = torch.cat([v_ctx, v_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim) + k = self.k_norm(k).transpose(1, 2) + v = v.transpose(1, 2) + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb(q, k, cos, sin) + if past_key_values is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs) + attn_fn: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + attn_output, attn_weights = attn_fn( + self, + q, + k, + v, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, q_len, -1) + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + +class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Qwen3DFlashAttention(config=config, layer_idx=layer_idx) + self.mlp = Qwen3MLP(config) + self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + target_hidden: Optional[torch.Tensor] = None, + hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + hidden_states=hidden_states, + target_hidden=target_hidden, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + )[0] + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + +def build_target_layer_ids(num_target_layers: int, num_draft_layers: int): + if num_draft_layers == 1: + return [(num_target_layers // 2)] + start = 1 + end = num_target_layers - 3 + span = end - start + target_layer_ids = [ + int(round(start + (i * span) / (num_draft_layers - 1))) + for i in range(num_draft_layers) + ] + return target_layer_ids + +def extract_context_feature( + hidden_states: list[torch.Tensor], + layer_ids: Optional[list[int]], +) -> torch.Tensor: + offset = 1 + selected_states = [] + for layer_id in layer_ids: + selected_states.append(hidden_states[layer_id + offset]) + target_hidden = torch.cat(selected_states, dim=-1) + return target_hidden + +class DFlashDraftModel(Qwen3PreTrainedModel): + config_class = Qwen3Config + _no_split_modules = ["Qwen3DFlashDecoderLayer"] + def __init__(self, config) -> None: + super().__init__(config) + self.config = config + self.layers = nn.ModuleList( + [Qwen3DFlashDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.target_layer_ids = build_target_layer_ids(config.num_target_layers, config.num_hidden_layers) + self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3RotaryEmbedding(config) + self.fc = nn.Linear(len(self.target_layer_ids) * config.hidden_size, config.hidden_size, bias=False) + self.hidden_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.block_size = config.block_size + self.post_init() + + def forward( + self, + position_ids: torch.LongTensor, + attention_mask: Optional[torch.Tensor] = None, + noise_embedding: Optional[torch.Tensor] = None, + target_hidden: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: bool = False, + **kwargs, + ) -> CausalLMOutputWithPast: + hidden_states = noise_embedding + target_hidden = self.hidden_norm(self.fc(target_hidden)) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for layer in self.layers: + hidden_states = layer( + hidden_states=hidden_states, + target_hidden=target_hidden, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + return self.norm(hidden_states) + + @torch.inference_mode() + def spec_generate( + self, + target: nn.Module, + input_ids: torch.LongTensor, + mask_token_id: int, + max_new_tokens: int, + stop_token_ids: list[int], + temperature: float, + ): + self.eval() + target.eval() + num_input_tokens = input_ids.shape[1] + max_length = num_input_tokens + max_new_tokens + + block_size = self.block_size + output_ids = torch.full( + (1, max_length + block_size), + mask_token_id, + dtype=torch.long, + device=target.device, + ) + position_ids = torch.arange(output_ids.shape[1], device=target.device).unsqueeze(0) + + past_key_values_target = DynamicCache() + past_key_values_draft = DynamicCache() + + # Prefill stage + output = target( + input_ids, + position_ids=position_ids[:, :num_input_tokens], + past_key_values=past_key_values_target, + use_cache=True, + logits_to_keep=1, + output_hidden_states=True, + ) + + output_ids[:, :num_input_tokens] = input_ids + output_ids[:, num_input_tokens:num_input_tokens+1] = sample(output.logits, temperature) + target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids) + + # Decode stage + acceptance_lengths = [] + start = input_ids.shape[1] + while start < max_length: + block_output_ids = output_ids[:, start : start + block_size].clone() + block_position_ids = position_ids[:, start : start + block_size] + noise_embedding = target.model.embed_tokens(block_output_ids) + draft_logits = target.lm_head(self( + target_hidden=target_hidden, + noise_embedding=noise_embedding, + position_ids=position_ids[:, past_key_values_draft.get_seq_length(): start + block_size], + past_key_values=past_key_values_draft, + use_cache=True, + is_causal=False, + )[:, -block_size+1:, :]) + past_key_values_draft.crop(start) + block_output_ids[:, 1:] = sample(draft_logits) + + output = target( + block_output_ids, + position_ids=block_position_ids, + past_key_values=past_key_values_target, + use_cache=True, + output_hidden_states=True, + ) + + posterior = sample(output.logits, temperature) + acceptance_length = (block_output_ids[:, 1:] == posterior[:, :-1]).cumprod(dim=1).sum(dim=1)[0].item() + output_ids[:, start : start + acceptance_length + 1] = block_output_ids[:, : acceptance_length + 1] + output_ids[:, start + acceptance_length + 1] = posterior[:, acceptance_length] + start += acceptance_length + 1 + past_key_values_target.crop(start) + target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids)[:, :acceptance_length + 1, :] + acceptance_lengths.append(acceptance_length+1) + if stop_token_ids is not None and any( + stop_token_id in output_ids[:, num_input_tokens:] for stop_token_id in stop_token_ids + ): + break + output_ids = output_ids[:, :max_length] + output_ids = output_ids[:, output_ids[0] != mask_token_id] + if stop_token_ids is not None: + stop_token_ids = torch.tensor(stop_token_ids, device=output_ids.device) + stop_token_indices = torch.isin(output_ids[0][num_input_tokens:], stop_token_ids).nonzero(as_tuple=True)[0] + if stop_token_indices.numel() > 0: + output_ids = output_ids[:, : num_input_tokens + stop_token_indices[0] + 1] + + return output_ids diff --git a/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/tokenwise_time.json b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/tokenwise_time.json new file mode 100644 index 0000000000000000000000000000000000000000..1a219fa903863033ff59e7d04d0534281da50af3 --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/tokenwise_time.json @@ -0,0 +1 @@ +{"bins": 200, "beta": 0.99, "updates": 29, "t": [0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0], "tau": [[0.0, 0.0396, 0.0807, 0.122, 0.1645, 0.2052, 0.2469, 0.2888, 0.3306, 0.373, 0.4158, 0.4606, 0.5054, 0.552, 0.6004, 0.6523, 0.7076, 0.7689, 0.8392, 0.9115, 1.0], [0.0, 0.0382, 0.0774, 0.1161, 0.1548, 0.1943, 0.2331, 0.2726, 0.3134, 0.3529, 0.3941, 0.4356, 0.4774, 0.5219, 0.5675, 0.6178, 0.6713, 0.732, 0.8069, 0.8926, 1.0], [0.0, 0.0367, 0.0739, 0.1119, 0.1496, 0.1883, 0.2257, 0.2634, 0.302, 0.3401, 0.379, 0.4188, 0.4598, 0.5019, 0.5464, 0.5938, 0.6457, 0.7057, 0.7813, 0.8759, 1.0], [0.0, 0.0361, 0.0727, 0.1095, 0.1465, 0.1838, 0.2211, 0.2588, 0.2964, 0.3341, 0.372, 0.411, 0.4515, 0.4924, 0.5355, 0.5814, 0.632, 0.6912, 0.7651, 0.8664, 1.0], [0.0, 0.0357, 0.0716, 0.1078, 0.1441, 0.1808, 0.2173, 0.2541, 0.2911, 0.3288, 0.3663, 0.4048, 0.4447, 0.4849, 0.5273, 0.5722, 0.6225, 0.6805, 0.7548, 0.8587, 1.0], [0.0, 0.0355, 0.0711, 0.1068, 0.143, 0.1795, 0.2154, 0.2518, 0.2886, 0.3262, 0.3633, 0.4012, 0.4404, 0.4805, 0.5228, 0.5678, 0.6176, 0.675, 0.7473, 0.8523, 1.0], [0.0, 0.0354, 0.0709, 0.1064, 0.1421, 0.1788, 0.2151, 0.2513, 0.2878, 0.325, 0.3619, 0.3996, 0.4385, 0.4778, 0.5196, 0.5638, 0.613, 0.6692, 0.74, 0.8461, 1.0], [0.0, 0.0353, 0.0709, 0.1063, 0.1421, 0.1781, 0.2141, 0.2501, 0.2862, 0.3232, 0.3603, 0.3979, 0.4365, 0.4756, 0.5172, 0.5613, 0.6094, 0.6651, 0.7356, 0.8419, 1.0], [0.0, 0.0353, 0.0707, 0.1059, 0.1415, 0.177, 0.2128, 0.2488, 0.2851, 0.3221, 0.3593, 0.3971, 0.4354, 0.4745, 0.5158, 0.56, 0.6077, 0.6629, 0.7323, 0.8393, 1.0], [0.0, 0.0353, 0.0705, 0.106, 0.1415, 0.1769, 0.2127, 0.2483, 0.2842, 0.3209, 0.3581, 0.3958, 0.4345, 0.473, 0.5142, 0.5578, 0.6056, 0.6596, 0.7293, 0.8359, 1.0], [0.0, 0.0353, 0.0705, 0.106, 0.1412, 0.1768, 0.2125, 0.2482, 0.2843, 0.321, 0.358, 0.3954, 0.4338, 0.4725, 0.513, 0.5565, 0.6041, 0.6578, 0.7265, 0.8336, 1.0], [0.0, 0.0353, 0.0707, 0.1064, 0.142, 0.1774, 0.2131, 0.2491, 0.2851, 0.3219, 0.3586, 0.3961, 0.4343, 0.4729, 0.5134, 0.5565, 0.6038, 0.658, 0.7261, 0.8327, 1.0], [0.0, 0.0357, 0.0713, 0.1068, 0.1427, 0.1782, 0.2141, 0.2501, 0.2864, 0.3231, 0.3605, 0.3978, 0.4359, 0.4746, 0.5154, 0.5591, 0.6066, 0.6602, 0.7286, 0.8344, 1.0], [0.0, 0.0358, 0.0717, 0.1074, 0.1438, 0.18, 0.2165, 0.2529, 0.2897, 0.327, 0.3645, 0.4026, 0.4412, 0.4805, 0.5217, 0.5655, 0.6126, 0.6669, 0.7343, 0.8381, 1.0], [0.0, 0.0372, 0.0744, 0.1116, 0.1488, 0.1862, 0.2237, 0.2615, 0.2991, 0.3375, 0.3761, 0.4151, 0.4546, 0.4952, 0.5374, 0.5824, 0.6307, 0.687, 0.7579, 0.8576, 1.0]], "loss_table": [[0.56088, 0.51558, 0.47125, 0.55303, 0.50507, 0.53824, 0.55543, 0.51585, 0.5447, 0.46988, 0.48713, 0.5043, 0.54745, 0.51792, 0.50553, 0.49522, 0.46846, 0.50874, 0.51544, 0.52513, 0.58333, 0.49041, 0.48576, 0.48477, 0.45076, 0.48379, 0.55411, 0.49368, 0.49259, 0.47939, 0.49061, 0.46604, 0.49517, 0.46604, 0.56361, 0.51253, 0.49062, 0.49547, 0.55671, 0.49089, 0.52201, 0.52091, 0.47561, 0.54581, 0.47084, 0.49916, 0.51017, 0.52189, 0.45984, 0.50827, 0.5119, 0.50995, 0.48959, 0.50931, 0.4923, 0.45744, 0.50718, 0.50554, 0.46289, 0.48852, 0.49966, 0.49434, 0.50693, 0.53308, 0.50433, 0.50986, 0.46111, 0.52461, 0.47881, 0.54803, 0.48718, 0.46517, 0.48453, 0.49171, 0.49162, 0.50611, 0.49551, 0.46613, 0.51336, 0.48207, 0.48288, 0.46742, 0.48628, 0.47021, 0.47866, 0.48778, 0.44934, 0.48692, 0.46111, 0.43962, 0.43403, 0.48642, 0.46903, 0.48561, 0.45442, 0.49051, 0.45499, 0.4871, 0.48842, 0.442, 0.4295, 0.42935, 0.43266, 0.44055, 0.47079, 0.46993, 0.43197, 0.46627, 0.44402, 0.44, 0.44536, 0.44512, 0.42112, 0.40952, 0.44025, 0.4679, 0.43347, 0.41875, 0.41953, 0.42335, 0.37395, 0.44334, 0.4165, 0.38566, 0.39681, 0.37053, 0.39293, 0.41946, 0.41616, 0.41201, 0.39121, 0.39031, 0.34707, 0.3994, 0.38, 0.40136, 0.34779, 0.36477, 0.4093, 0.35792, 0.38638, 0.35186, 0.37903, 0.38355, 0.35282, 0.35586, 0.32737, 0.34689, 0.3105, 0.33526, 0.31963, 0.32623, 0.30799, 0.33363, 0.30344, 0.30949, 0.32834, 0.32899, 0.29014, 0.27143, 0.26848, 0.24515, 0.32626, 0.30014, 0.30361, 0.28792, 0.27329, 0.3174, 0.29754, 0.27771, 0.25178, 0.28983, 0.27826, 0.33397, 0.28043, 0.27025, 0.33233, 0.2596, 0.27997, 0.31038, 0.31204, 0.27314, 0.25264, 0.2712, 0.32573, 0.27814, 0.21979, 0.28643, 0.27291, 0.2474, 0.25133, 0.20983, 0.23553, 0.24367, 0.20727, 0.21999, 0.21533, 0.17414, 0.20101, 0.13969], [1.02762, 0.99538, 0.88432, 1.03223, 0.99607, 0.94777, 1.01154, 0.93535, 0.97038, 0.95691, 0.98914, 0.92718, 0.9606, 0.97058, 0.92196, 0.93416, 0.92885, 0.93895, 0.97835, 0.99102, 1.04719, 0.9829, 0.9134, 0.96962, 0.99004, 0.9783, 1.07615, 0.99307, 0.92447, 0.90467, 0.89409, 0.9108, 0.95659, 0.87792, 0.97548, 0.96517, 0.9452, 0.95927, 1.00171, 0.99361, 0.99984, 0.99025, 0.9442, 0.97116, 0.90401, 0.96519, 0.9426, 0.92727, 0.92647, 0.97093, 0.9458, 0.95741, 0.95577, 0.93808, 0.94552, 0.90474, 0.93727, 0.92243, 0.89924, 0.89257, 0.95599, 0.88696, 0.94607, 0.93743, 0.95191, 0.95026, 0.91248, 0.98515, 0.9636, 0.91119, 0.97326, 0.89274, 0.90524, 0.92688, 0.88133, 0.9491, 0.87298, 0.89009, 0.92506, 0.96889, 0.91295, 0.87118, 0.88589, 0.89584, 0.91134, 0.88502, 0.88745, 0.92258, 0.90019, 0.92157, 0.83013, 0.93339, 0.89736, 0.90471, 0.86812, 0.85899, 0.87822, 0.82619, 0.87228, 0.82634, 0.82514, 0.81643, 0.84316, 0.82564, 0.86575, 0.87875, 0.79413, 0.75266, 0.78757, 0.88704, 0.82372, 0.84408, 0.78424, 0.78401, 0.78807, 0.77592, 0.7541, 0.76547, 0.73941, 0.73318, 0.6793, 0.74895, 0.71661, 0.6989, 0.74665, 0.72632, 0.6692, 0.72292, 0.72928, 0.71417, 0.70138, 0.71887, 0.62195, 0.6543, 0.67166, 0.68749, 0.57142, 0.64435, 0.6259, 0.60847, 0.61048, 0.58654, 0.6167, 0.59466, 0.62625, 0.57759, 0.58388, 0.55195, 0.50945, 0.53128, 0.53833, 0.52966, 0.53456, 0.49407, 0.48102, 0.4915, 0.48779, 0.4682, 0.45283, 0.47576, 0.42874, 0.42646, 0.46813, 0.45699, 0.5021, 0.48529, 0.43728, 0.43921, 0.4256, 0.43359, 0.41289, 0.45334, 0.38461, 0.45049, 0.42304, 0.44604, 0.44099, 0.36123, 0.38382, 0.39788, 0.47385, 0.38239, 0.40329, 0.40414, 0.40447, 0.39779, 0.36871, 0.37758, 0.38914, 0.33856, 0.36904, 0.29247, 0.34438, 0.36225, 0.32639, 0.27743, 0.31428, 0.24372, 0.25904, 0.17693], [1.39247, 1.38406, 1.25961, 1.42022, 1.41178, 1.40185, 1.41545, 1.30526, 1.44834, 1.39428, 1.35023, 1.30685, 1.37198, 1.37245, 1.31612, 1.2943, 1.33304, 1.32829, 1.32196, 1.38355, 1.38361, 1.3113, 1.28669, 1.3425, 1.35155, 1.34296, 1.43906, 1.35421, 1.28963, 1.31364, 1.31461, 1.24478, 1.32654, 1.26231, 1.36446, 1.32051, 1.32026, 1.32412, 1.40552, 1.39609, 1.36348, 1.37519, 1.2652, 1.35077, 1.3234, 1.3659, 1.34366, 1.35858, 1.2994, 1.35514, 1.30413, 1.34826, 1.39355, 1.28966, 1.34892, 1.23785, 1.34728, 1.31551, 1.2906, 1.28194, 1.4105, 1.2745, 1.34553, 1.32836, 1.2922, 1.29844, 1.33686, 1.39398, 1.32354, 1.27727, 1.34379, 1.27164, 1.33153, 1.31322, 1.26039, 1.30327, 1.27951, 1.27229, 1.3401, 1.29432, 1.23657, 1.21544, 1.22778, 1.28945, 1.31309, 1.22509, 1.19337, 1.25486, 1.21533, 1.2803, 1.17303, 1.21369, 1.21296, 1.2955, 1.19461, 1.15021, 1.22812, 1.23281, 1.1915, 1.13653, 1.17004, 1.13911, 1.18272, 1.08145, 1.21147, 1.17842, 1.0878, 1.07052, 1.13294, 1.16885, 1.09225, 1.07897, 1.11282, 1.01316, 1.04632, 1.10286, 1.05417, 1.03913, 0.97601, 1.0391, 0.97449, 1.00173, 0.98357, 0.9542, 0.98566, 1.00895, 0.92706, 0.94812, 0.94763, 0.9082, 0.90004, 0.89454, 0.86265, 0.85017, 0.85458, 0.85139, 0.84949, 0.79366, 0.79639, 0.79003, 0.81175, 0.71356, 0.77763, 0.78755, 0.70645, 0.74831, 0.70985, 0.6955, 0.6052, 0.66888, 0.57857, 0.66656, 0.67551, 0.57231, 0.60089, 0.57981, 0.55275, 0.56903, 0.55284, 0.56615, 0.55746, 0.50724, 0.57554, 0.54949, 0.58663, 0.53995, 0.51674, 0.53292, 0.57764, 0.50524, 0.4991, 0.5182, 0.49202, 0.4968, 0.47355, 0.54508, 0.48482, 0.48534, 0.44467, 0.49241, 0.49128, 0.43389, 0.45196, 0.45235, 0.43132, 0.45529, 0.41751, 0.47448, 0.43961, 0.39298, 0.38583, 0.38189, 0.42117, 0.38821, 0.35213, 0.32267, 0.32789, 0.26984, 0.30796, 0.18058], [1.73223, 1.7755, 1.60913, 1.79926, 1.76108, 1.64749, 1.69944, 1.67286, 1.74003, 1.72781, 1.66074, 1.6537, 1.71424, 1.68032, 1.61124, 1.67927, 1.65061, 1.68088, 1.64426, 1.68928, 1.78297, 1.67711, 1.60839, 1.68218, 1.69622, 1.62467, 1.78005, 1.6707, 1.65837, 1.66272, 1.66166, 1.68785, 1.62045, 1.6244, 1.71604, 1.66162, 1.60908, 1.6545, 1.71802, 1.63318, 1.71253, 1.63038, 1.63466, 1.66389, 1.63579, 1.72739, 1.62524, 1.63875, 1.67306, 1.63177, 1.58273, 1.62124, 1.66628, 1.59659, 1.68948, 1.63038, 1.66174, 1.6504, 1.63493, 1.539, 1.68561, 1.62392, 1.6937, 1.6905, 1.60508, 1.63524, 1.63098, 1.68033, 1.65764, 1.62827, 1.66204, 1.61378, 1.62376, 1.6037, 1.53998, 1.6189, 1.52575, 1.61762, 1.6571, 1.66676, 1.60848, 1.46078, 1.54165, 1.57567, 1.61027, 1.49714, 1.44721, 1.56218, 1.49093, 1.52199, 1.49581, 1.51234, 1.52039, 1.56436, 1.46232, 1.46908, 1.52703, 1.52481, 1.50997, 1.45905, 1.46133, 1.41879, 1.41737, 1.39525, 1.45727, 1.46058, 1.38433, 1.38428, 1.43179, 1.39322, 1.35345, 1.32672, 1.32276, 1.33616, 1.29273, 1.32114, 1.29924, 1.2585, 1.244, 1.25159, 1.23503, 1.23348, 1.24236, 1.12755, 1.21586, 1.17768, 1.14899, 1.20921, 1.13216, 1.0668, 1.00905, 1.03041, 0.99304, 1.00101, 1.01254, 1.02046, 0.98293, 0.98276, 0.9988, 0.9327, 0.91558, 0.94347, 0.91974, 0.91463, 0.80197, 0.85995, 0.78762, 0.81561, 0.70711, 0.74381, 0.72987, 0.74762, 0.78596, 0.61075, 0.68011, 0.68534, 0.65152, 0.70369, 0.64562, 0.6411, 0.62759, 0.56633, 0.66411, 0.58995, 0.60712, 0.6005, 0.57258, 0.54448, 0.64266, 0.55625, 0.54688, 0.5742, 0.49621, 0.62562, 0.5693, 0.54512, 0.51883, 0.52462, 0.49653, 0.51221, 0.56307, 0.49268, 0.50355, 0.54227, 0.514, 0.51658, 0.48565, 0.53674, 0.50523, 0.42903, 0.45687, 0.41596, 0.46201, 0.44731, 0.37892, 0.38808, 0.31884, 0.27808, 0.3284, 0.20143], [1.97126, 2.08037, 1.90608, 2.0225, 2.07594, 1.92145, 1.96177, 1.98986, 2.00115, 2.03013, 2.03682, 1.95524, 1.97639, 1.91878, 1.8989, 1.99178, 1.93718, 1.97153, 1.95258, 1.98456, 2.00565, 1.91504, 1.91812, 1.96398, 1.92278, 1.94249, 2.0507, 1.98532, 1.96421, 1.95409, 1.91143, 1.91287, 1.94543, 1.93829, 1.97525, 1.94248, 1.97382, 1.92359, 1.95789, 1.91097, 2.05972, 1.92409, 1.90063, 1.93666, 1.98724, 1.97588, 1.9345, 1.87683, 1.96674, 1.90331, 1.84661, 1.89825, 1.97702, 1.93353, 1.91475, 1.87314, 1.9655, 1.9246, 1.82814, 1.83979, 1.90342, 1.84831, 1.94839, 1.90732, 1.91583, 1.91716, 1.85779, 1.96533, 1.93294, 1.91351, 1.89157, 1.84167, 1.88035, 1.83268, 1.7964, 1.8815, 1.85157, 1.84454, 1.86801, 1.88055, 1.80864, 1.70796, 1.84623, 1.84818, 1.87898, 1.76972, 1.72057, 1.81081, 1.6984, 1.7711, 1.80347, 1.76713, 1.75191, 1.83051, 1.73397, 1.76624, 1.7259, 1.7206, 1.73823, 1.62979, 1.59856, 1.66426, 1.7394, 1.61395, 1.68658, 1.72749, 1.63855, 1.61201, 1.62635, 1.64742, 1.60493, 1.54626, 1.46957, 1.48811, 1.48467, 1.48309, 1.44593, 1.41797, 1.45297, 1.38391, 1.38196, 1.42553, 1.38703, 1.35507, 1.36824, 1.32788, 1.26502, 1.33334, 1.25181, 1.2521, 1.19892, 1.22811, 1.10653, 1.16323, 1.14006, 1.15436, 1.0692, 1.04438, 1.14141, 1.06509, 1.06131, 0.99856, 0.98601, 0.97514, 0.88663, 0.93474, 0.89377, 0.9366, 0.79659, 0.80256, 0.77371, 0.86765, 0.83331, 0.76765, 0.72854, 0.76583, 0.67592, 0.71234, 0.70147, 0.72735, 0.63839, 0.62299, 0.67196, 0.64772, 0.68219, 0.64364, 0.63281, 0.59526, 0.61076, 0.60775, 0.60326, 0.62464, 0.5648, 0.68379, 0.61686, 0.61613, 0.59626, 0.55689, 0.50304, 0.54579, 0.63629, 0.5264, 0.50529, 0.58868, 0.50082, 0.55016, 0.48655, 0.57274, 0.52673, 0.49394, 0.50557, 0.4614, 0.45599, 0.48904, 0.41707, 0.42265, 0.38373, 0.34169, 0.31872, 0.19825], [2.29948, 2.32162, 2.19231, 2.26537, 2.26241, 2.13559, 2.20833, 2.19048, 2.32071, 2.27736, 2.26993, 2.18542, 2.22269, 2.20385, 2.17519, 2.25906, 2.21508, 2.21766, 2.21907, 2.22475, 2.28629, 2.17539, 2.13424, 2.21827, 2.23114, 2.16394, 2.31205, 2.1835, 2.13948, 2.16864, 2.18447, 2.20678, 2.1782, 2.17311, 2.21006, 2.15668, 2.20512, 2.27438, 2.22203, 2.22484, 2.31204, 2.11205, 2.16434, 2.21357, 2.17164, 2.27059, 2.20433, 2.10947, 2.16184, 2.18248, 2.09832, 2.22668, 2.15529, 2.12148, 2.16426, 2.20675, 2.15353, 2.14081, 2.13692, 2.04048, 2.20354, 2.16712, 2.13363, 2.04236, 2.09574, 2.16475, 2.10062, 2.15669, 2.19656, 2.15811, 2.14745, 2.08735, 2.13835, 2.09538, 2.05957, 2.17074, 2.1248, 2.10821, 2.11247, 2.01868, 2.09146, 1.94866, 2.03071, 2.08639, 2.06046, 1.95902, 1.99248, 2.05275, 1.98023, 2.05391, 1.96141, 1.9765, 1.9662, 2.05133, 1.94789, 1.95636, 1.9139, 1.94715, 1.877, 1.88535, 1.85557, 1.83223, 1.88812, 1.86214, 1.86087, 1.90436, 1.76197, 1.8038, 1.78573, 1.77042, 1.76762, 1.69047, 1.64532, 1.69977, 1.65023, 1.70298, 1.59929, 1.56564, 1.60484, 1.54357, 1.57212, 1.63075, 1.52134, 1.48676, 1.55621, 1.52791, 1.40003, 1.45327, 1.44679, 1.30968, 1.36206, 1.30888, 1.28594, 1.2497, 1.30543, 1.25327, 1.24612, 1.20588, 1.28924, 1.17645, 1.20175, 1.14743, 1.0901, 1.05679, 0.99184, 1.01213, 0.96798, 0.99818, 0.85366, 0.92132, 0.84789, 0.89244, 0.91129, 0.82771, 0.77691, 0.84395, 0.77573, 0.79755, 0.72172, 0.74945, 0.73663, 0.66551, 0.74837, 0.71232, 0.6938, 0.70028, 0.66079, 0.71968, 0.70176, 0.63633, 0.6573, 0.65824, 0.6529, 0.68353, 0.66064, 0.6619, 0.63388, 0.56119, 0.56407, 0.58616, 0.61582, 0.54014, 0.63408, 0.58378, 0.55919, 0.57138, 0.54867, 0.62126, 0.50698, 0.53344, 0.55009, 0.49392, 0.48295, 0.48856, 0.42881, 0.41303, 0.41505, 0.31643, 0.34947, 0.2366], [2.47359, 2.51735, 2.42206, 2.5459, 2.5669, 2.36877, 2.39208, 2.44683, 2.54713, 2.5176, 2.49642, 2.39543, 2.41564, 2.4165, 2.39209, 2.46109, 2.40138, 2.5069, 2.42709, 2.48118, 2.53097, 2.44242, 2.43031, 2.49249, 2.46437, 2.39914, 2.56485, 2.35905, 2.3801, 2.40193, 2.42502, 2.31485, 2.41298, 2.36776, 2.44866, 2.29361, 2.41199, 2.42504, 2.4588, 2.40675, 2.44178, 2.32894, 2.37266, 2.40121, 2.40862, 2.48079, 2.4465, 2.32679, 2.39651, 2.42232, 2.37273, 2.39648, 2.37127, 2.44724, 2.35801, 2.41787, 2.40955, 2.34589, 2.36929, 2.30641, 2.46273, 2.3193, 2.37259, 2.35799, 2.28968, 2.35178, 2.35003, 2.37437, 2.36346, 2.42206, 2.35112, 2.3048, 2.43637, 2.32979, 2.31766, 2.3895, 2.25347, 2.3401, 2.31918, 2.17544, 2.32717, 2.19994, 2.25261, 2.29007, 2.24908, 2.16508, 2.23141, 2.30002, 2.24234, 2.23647, 2.19908, 2.18255, 2.20361, 2.27208, 2.21676, 2.15747, 2.19303, 2.11893, 2.07926, 2.10254, 2.06002, 2.01838, 2.04424, 2.07531, 2.0396, 2.10542, 1.9939, 1.98383, 1.96747, 1.94138, 1.95478, 1.87451, 1.89261, 1.85529, 1.78825, 1.8502, 1.77745, 1.72942, 1.74563, 1.75454, 1.74762, 1.73999, 1.72329, 1.68798, 1.70473, 1.63008, 1.60462, 1.5829, 1.55828, 1.49039, 1.48367, 1.47681, 1.39789, 1.38646, 1.43664, 1.37043, 1.37939, 1.28682, 1.34367, 1.32691, 1.30066, 1.21403, 1.135, 1.15784, 1.07083, 1.1534, 1.05769, 1.01529, 0.9416, 1.02398, 0.96144, 0.88267, 0.91814, 0.92179, 0.88818, 0.91228, 0.8559, 0.80119, 0.75222, 0.78999, 0.7796, 0.79787, 0.82216, 0.73766, 0.73398, 0.68462, 0.70436, 0.71014, 0.70556, 0.65947, 0.70401, 0.67912, 0.65875, 0.68077, 0.69917, 0.68804, 0.69367, 0.59218, 0.61218, 0.66388, 0.66145, 0.56756, 0.64432, 0.58344, 0.63211, 0.55701, 0.57902, 0.61186, 0.51541, 0.58183, 0.5992, 0.55129, 0.53593, 0.50489, 0.44667, 0.42801, 0.37137, 0.31359, 0.36652, 0.22903], [2.62318, 2.68436, 2.63953, 2.72575, 2.72539, 2.60525, 2.54459, 2.60722, 2.69788, 2.66432, 2.64592, 2.60388, 2.63215, 2.53109, 2.62323, 2.63314, 2.65937, 2.65652, 2.61811, 2.65688, 2.68228, 2.59887, 2.59056, 2.69192, 2.60097, 2.60316, 2.67847, 2.56658, 2.5972, 2.54771, 2.69296, 2.55566, 2.67831, 2.57165, 2.58588, 2.53559, 2.57239, 2.61148, 2.66139, 2.59256, 2.63844, 2.55564, 2.57193, 2.58441, 2.62833, 2.66928, 2.56815, 2.58109, 2.60341, 2.56454, 2.63079, 2.59884, 2.66072, 2.57815, 2.57533, 2.56039, 2.52987, 2.56814, 2.57346, 2.49797, 2.61077, 2.45763, 2.52386, 2.49081, 2.53966, 2.44319, 2.51663, 2.57175, 2.56543, 2.47786, 2.52964, 2.51362, 2.59191, 2.51921, 2.45442, 2.50109, 2.49034, 2.4619, 2.47715, 2.36941, 2.48532, 2.41363, 2.43114, 2.51283, 2.42021, 2.32243, 2.40102, 2.45977, 2.38968, 2.41212, 2.46207, 2.37965, 2.37786, 2.38334, 2.33735, 2.31578, 2.26334, 2.27127, 2.27081, 2.27294, 2.1668, 2.18305, 2.2522, 2.24842, 2.22891, 2.20283, 2.12711, 2.06252, 2.10179, 2.10422, 2.03152, 2.06164, 2.05227, 2.00477, 2.01331, 2.01061, 1.94977, 1.87082, 1.95185, 1.87767, 1.87069, 1.85476, 1.78504, 1.79267, 1.78988, 1.74247, 1.73071, 1.68164, 1.6433, 1.64647, 1.53901, 1.57002, 1.539, 1.44472, 1.51117, 1.49991, 1.4721, 1.41414, 1.4426, 1.29382, 1.3529, 1.313, 1.25755, 1.18903, 1.12549, 1.18881, 1.115, 1.11362, 0.98455, 1.062, 0.97614, 1.01485, 0.94196, 0.97848, 0.9171, 0.95446, 0.9208, 0.88427, 0.8355, 0.80805, 0.83618, 0.77923, 0.74079, 0.79172, 0.79593, 0.75683, 0.69412, 0.75011, 0.79956, 0.6692, 0.68808, 0.72722, 0.68257, 0.74523, 0.70298, 0.66785, 0.68852, 0.65846, 0.64833, 0.71774, 0.65158, 0.62951, 0.66502, 0.60977, 0.62296, 0.60892, 0.58539, 0.62181, 0.61424, 0.63007, 0.61194, 0.51127, 0.50514, 0.49971, 0.45259, 0.41851, 0.3856, 0.38595, 0.37841, 0.23612], [2.82419, 2.92267, 2.82569, 2.92023, 2.86845, 2.82141, 2.80209, 2.84906, 2.92991, 2.86659, 2.83081, 2.81449, 2.84483, 2.80286, 2.83685, 2.89596, 2.88936, 2.86864, 2.78336, 2.89779, 2.87942, 2.82319, 2.74958, 2.89023, 2.84567, 2.79206, 2.93384, 2.81109, 2.83433, 2.84752, 2.89032, 2.84853, 2.78525, 2.84456, 2.84237, 2.76241, 2.84098, 2.75858, 2.88887, 2.81955, 2.87501, 2.76508, 2.83625, 2.80329, 2.83168, 2.81657, 2.75809, 2.79208, 2.8253, 2.75134, 2.82072, 2.83687, 2.80838, 2.74163, 2.76983, 2.78096, 2.70803, 2.83452, 2.73566, 2.70242, 2.82788, 2.63962, 2.77372, 2.61207, 2.63566, 2.65495, 2.71234, 2.78421, 2.77395, 2.70352, 2.69608, 2.67341, 2.69233, 2.71602, 2.55893, 2.72036, 2.66004, 2.70568, 2.63769, 2.67469, 2.61028, 2.56249, 2.63501, 2.76148, 2.63692, 2.51733, 2.63762, 2.63352, 2.50485, 2.59365, 2.64709, 2.6075, 2.58196, 2.55877, 2.54121, 2.54955, 2.4715, 2.46349, 2.38599, 2.47717, 2.34273, 2.38537, 2.42968, 2.37872, 2.46197, 2.3476, 2.34792, 2.22425, 2.22211, 2.21652, 2.19202, 2.15866, 2.17138, 2.1962, 2.20693, 2.11529, 2.14265, 2.04314, 2.07428, 2.09459, 2.03244, 2.04183, 1.99073, 1.91163, 2.00481, 1.79645, 1.80112, 1.81571, 1.7579, 1.81156, 1.71715, 1.69223, 1.60108, 1.55026, 1.66163, 1.54717, 1.6153, 1.55495, 1.55324, 1.44619, 1.41721, 1.45099, 1.29747, 1.3363, 1.27268, 1.24945, 1.18074, 1.20979, 1.14556, 1.12436, 0.99397, 1.11128, 0.98087, 1.06874, 0.98153, 1.02682, 0.95536, 0.91077, 0.88113, 0.89216, 0.82765, 0.83009, 0.77034, 0.79597, 0.81433, 0.81407, 0.72453, 0.78853, 0.7526, 0.75931, 0.79678, 0.7755, 0.73067, 0.70257, 0.67445, 0.69331, 0.7876, 0.75945, 0.6748, 0.74804, 0.75049, 0.67241, 0.6849, 0.63639, 0.67249, 0.66684, 0.65885, 0.65661, 0.63446, 0.62242, 0.55868, 0.56933, 0.47697, 0.53679, 0.4607, 0.48541, 0.42857, 0.37027, 0.39667, 0.2613], [2.96078, 3.08528, 2.98202, 3.11185, 3.07657, 2.9585, 2.93471, 3.06781, 3.08004, 2.99539, 3.03066, 2.94857, 3.09964, 2.99414, 2.93686, 3.02473, 2.98893, 3.03495, 2.96796, 3.07862, 2.98506, 2.96166, 2.9367, 3.09086, 3.05165, 2.87443, 3.05996, 3.03893, 3.07283, 3.0431, 3.06567, 2.97463, 2.98853, 3.02234, 2.97629, 2.9424, 2.99425, 2.95168, 2.95752, 3.0136, 3.04441, 2.94657, 2.92354, 3.00422, 3.00517, 3.00526, 3.02892, 2.94892, 3.03256, 2.94899, 3.03226, 2.99525, 2.99134, 2.97932, 2.85408, 2.98672, 2.90411, 2.94759, 2.84602, 2.92846, 3.05468, 2.88836, 2.83443, 2.88146, 2.84313, 2.81507, 2.86537, 2.88481, 2.91202, 2.86346, 2.87113, 2.81804, 2.90101, 2.83194, 2.76212, 2.85033, 2.81501, 2.83308, 2.84575, 2.72701, 2.73601, 2.68478, 2.80374, 2.84256, 2.78915, 2.71373, 2.74311, 2.73093, 2.71021, 2.79797, 2.7897, 2.76234, 2.75704, 2.84562, 2.73079, 2.60782, 2.64465, 2.57399, 2.6244, 2.60075, 2.52205, 2.53391, 2.49652, 2.5121, 2.51621, 2.49046, 2.53721, 2.44207, 2.44807, 2.31918, 2.37956, 2.31683, 2.36704, 2.30048, 2.22249, 2.22051, 2.22153, 2.19143, 2.21278, 2.16742, 2.13992, 2.21071, 2.0473, 2.02574, 2.11375, 1.97923, 1.91556, 1.99928, 1.84571, 1.95333, 1.76351, 1.84303, 1.68061, 1.63882, 1.72266, 1.68518, 1.62932, 1.60295, 1.54238, 1.48728, 1.57523, 1.45951, 1.37566, 1.31152, 1.36451, 1.28431, 1.21921, 1.30121, 1.20177, 1.15255, 1.06057, 1.11965, 1.07252, 1.11186, 1.04809, 0.99592, 1.02337, 0.97258, 0.85537, 0.90561, 0.92244, 0.89452, 0.87062, 0.87724, 0.79011, 0.81471, 0.77344, 0.89019, 0.84675, 0.7863, 0.79655, 0.79136, 0.78903, 0.69109, 0.79933, 0.76122, 0.74675, 0.74444, 0.70679, 0.77344, 0.75115, 0.64208, 0.71032, 0.65917, 0.60391, 0.70882, 0.61943, 0.7233, 0.60813, 0.57173, 0.63867, 0.55988, 0.58275, 0.53571, 0.50102, 0.47117, 0.42009, 0.38298, 0.43448, 0.23533], [3.19289, 3.20332, 3.13118, 3.29195, 3.21025, 3.11974, 3.09509, 3.24822, 3.29626, 3.14663, 3.14776, 3.13062, 3.24695, 3.10309, 3.07283, 3.12355, 3.19411, 3.19611, 3.10147, 3.24791, 3.15023, 3.20357, 3.08551, 3.2674, 3.22793, 3.09495, 3.25732, 3.19573, 3.11871, 3.15339, 3.18287, 3.16111, 3.13278, 3.18914, 3.09748, 3.11531, 3.18987, 3.13633, 3.13197, 3.11869, 3.23783, 3.08722, 3.10859, 3.11983, 3.13705, 3.20026, 3.13958, 3.10017, 3.18473, 3.0822, 3.16972, 3.11158, 3.06028, 3.1369, 3.07121, 3.16635, 3.06199, 3.10235, 3.0049, 3.07904, 3.17164, 2.96408, 3.02321, 3.02179, 3.06421, 2.974, 3.07155, 3.04761, 3.01762, 3.04448, 3.06716, 2.94967, 3.13895, 3.03683, 2.91446, 2.99686, 2.99002, 2.90381, 2.99524, 2.92879, 2.93743, 2.90347, 2.95506, 2.98542, 2.95227, 2.85467, 2.85478, 2.90819, 2.81035, 2.893, 3.00207, 2.86355, 2.92082, 2.92556, 2.80703, 2.80813, 2.81539, 2.85791, 2.81084, 2.85916, 2.66897, 2.62131, 2.66889, 2.60001, 2.72496, 2.65167, 2.69877, 2.55645, 2.51446, 2.41763, 2.46838, 2.43612, 2.48679, 2.4314, 2.40386, 2.38831, 2.31003, 2.29669, 2.24604, 2.33123, 2.24731, 2.26932, 2.20914, 2.15656, 2.27917, 1.99678, 2.00843, 2.07626, 2.01981, 2.02367, 1.93297, 1.86546, 1.8737, 1.74425, 1.77206, 1.72595, 1.6856, 1.68617, 1.61236, 1.59955, 1.65111, 1.5468, 1.49649, 1.39286, 1.40001, 1.38538, 1.27518, 1.29637, 1.24033, 1.28538, 1.13003, 1.0992, 1.08888, 1.1286, 1.13199, 1.04911, 1.00675, 1.01923, 0.93289, 0.91551, 0.91603, 0.92905, 0.87201, 0.85094, 0.8502, 0.86083, 0.81319, 0.8325, 0.85632, 0.78461, 0.80661, 0.78982, 0.76658, 0.74242, 0.77452, 0.86511, 0.7526, 0.75133, 0.68739, 0.75834, 0.83831, 0.7594, 0.75613, 0.75215, 0.68396, 0.75373, 0.66607, 0.66716, 0.64485, 0.57796, 0.61682, 0.60869, 0.64596, 0.59495, 0.51783, 0.48442, 0.41417, 0.36773, 0.44011, 0.23689], [3.37818, 3.4295, 3.32981, 3.39272, 3.38651, 3.268, 3.23625, 3.35117, 3.40409, 3.29405, 3.30436, 3.31998, 3.36086, 3.2574, 3.28336, 3.27311, 3.29237, 3.35928, 3.26916, 3.31437, 3.36866, 3.31251, 3.22481, 3.38328, 3.33019, 3.24876, 3.37842, 3.30472, 3.27889, 3.38972, 3.3586, 3.33149, 3.36475, 3.31092, 3.32532, 3.27716, 3.35026, 3.29164, 3.33912, 3.28857, 3.35237, 3.2818, 3.22231, 3.21942, 3.29937, 3.38346, 3.21985, 3.22702, 3.30349, 3.2728, 3.30489, 3.28141, 3.35561, 3.22667, 3.21544, 3.38781, 3.20666, 3.21179, 3.19872, 3.15896, 3.3226, 3.14365, 3.22858, 3.14938, 3.22504, 3.19481, 3.26758, 3.2346, 3.21266, 3.24453, 3.19085, 3.12807, 3.25126, 3.14237, 3.10679, 3.15485, 3.0806, 3.19298, 3.12594, 3.11598, 3.15629, 3.00586, 3.08751, 3.07022, 3.18309, 3.004, 3.05763, 3.14507, 2.98242, 3.15756, 3.10434, 2.99123, 3.04945, 3.07908, 2.89691, 2.97349, 3.00445, 2.96261, 2.93783, 2.9976, 2.81459, 2.76322, 2.87161, 2.76498, 2.91464, 2.80025, 2.78643, 2.73055, 2.67852, 2.55905, 2.64193, 2.60446, 2.62491, 2.50699, 2.56031, 2.55452, 2.4532, 2.40654, 2.46226, 2.42591, 2.35981, 2.35727, 2.21805, 2.30942, 2.39933, 2.10221, 2.11791, 2.13281, 2.08731, 2.05424, 2.04794, 1.99453, 1.95098, 1.80172, 1.9307, 1.8774, 1.85333, 1.80035, 1.77942, 1.7042, 1.64524, 1.5701, 1.54015, 1.50618, 1.49893, 1.45109, 1.32074, 1.42309, 1.25068, 1.28895, 1.19161, 1.25326, 1.1978, 1.16706, 1.22731, 1.09005, 1.03965, 1.09228, 0.96609, 0.96456, 0.97247, 0.9117, 0.95574, 0.89686, 0.89769, 0.92322, 0.81357, 0.91154, 0.89964, 0.81597, 0.79613, 0.84202, 0.80943, 0.77262, 0.90455, 0.83815, 0.81172, 0.79289, 0.728, 0.76648, 0.89912, 0.78447, 0.76121, 0.77736, 0.77273, 0.6895, 0.67911, 0.74645, 0.68729, 0.64679, 0.65137, 0.60097, 0.60363, 0.58089, 0.56474, 0.51068, 0.47838, 0.40479, 0.41969, 0.27905], [3.45301, 3.54717, 3.4596, 3.47587, 3.54046, 3.40331, 3.38482, 3.51748, 3.54336, 3.46751, 3.45968, 3.4795, 3.50039, 3.41229, 3.39587, 3.42847, 3.43456, 3.57913, 3.52596, 3.48395, 3.57092, 3.43764, 3.40874, 3.53676, 3.40964, 3.39528, 3.54721, 3.4197, 3.42946, 3.47917, 3.52677, 3.46229, 3.51521, 3.4927, 3.46864, 3.45446, 3.5334, 3.43759, 3.46917, 3.42714, 3.44916, 3.39819, 3.42604, 3.38423, 3.45926, 3.49558, 3.4668, 3.36662, 3.4788, 3.42074, 3.42612, 3.36835, 3.48685, 3.41875, 3.31003, 3.46762, 3.39454, 3.48105, 3.34501, 3.32963, 3.43769, 3.2927, 3.38675, 3.37021, 3.31648, 3.35511, 3.31156, 3.33986, 3.38425, 3.27367, 3.28704, 3.24931, 3.36041, 3.30686, 3.28591, 3.33311, 3.37636, 3.32563, 3.28749, 3.27392, 3.31845, 3.18109, 3.29035, 3.2906, 3.24871, 3.19364, 3.19441, 3.25998, 3.25125, 3.26543, 3.1979, 3.13131, 3.18998, 3.21753, 3.0947, 3.04775, 3.07987, 3.1097, 3.07547, 3.09425, 2.96571, 2.90049, 2.99277, 2.88027, 2.94319, 2.93843, 2.92938, 2.74973, 2.85985, 2.75994, 2.7209, 2.72269, 2.79425, 2.69455, 2.65603, 2.61492, 2.6297, 2.55784, 2.53342, 2.46743, 2.47238, 2.51682, 2.39722, 2.39615, 2.483, 2.25951, 2.32966, 2.39053, 2.17893, 2.1878, 2.1899, 2.18898, 2.0122, 1.9554, 2.00041, 2.0044, 1.93296, 1.87236, 1.87714, 1.79056, 1.77574, 1.66991, 1.62939, 1.64213, 1.56681, 1.54654, 1.53242, 1.50248, 1.42946, 1.33256, 1.28468, 1.27894, 1.28488, 1.24691, 1.27329, 1.18784, 1.10191, 1.11608, 1.06125, 1.04563, 1.01995, 0.99509, 1.02159, 0.91851, 0.92137, 0.94517, 0.93674, 0.94291, 0.86937, 0.92664, 0.87062, 0.86865, 0.84939, 0.89536, 0.96442, 0.83353, 0.86622, 0.88632, 0.85741, 0.83835, 0.897, 0.81296, 0.78341, 0.80402, 0.79045, 0.77768, 0.68289, 0.75662, 0.74469, 0.6454, 0.70259, 0.62594, 0.66973, 0.63104, 0.57764, 0.54656, 0.5149, 0.45211, 0.43891, 0.31838], [3.61713, 3.66843, 3.57499, 3.6992, 3.66139, 3.63133, 3.60992, 3.59796, 3.71214, 3.60397, 3.6604, 3.60093, 3.66519, 3.5744, 3.59816, 3.58898, 3.68096, 3.70525, 3.65361, 3.61855, 3.60859, 3.64611, 3.51597, 3.69308, 3.5018, 3.55028, 3.65746, 3.55612, 3.56723, 3.56024, 3.64178, 3.5331, 3.61779, 3.65581, 3.63479, 3.51192, 3.6095, 3.55331, 3.58343, 3.56697, 3.55165, 3.54466, 3.56945, 3.56452, 3.56764, 3.60176, 3.60566, 3.50674, 3.63727, 3.51895, 3.53621, 3.46425, 3.56547, 3.57512, 3.48257, 3.62521, 3.50364, 3.58131, 3.50373, 3.54714, 3.5253, 3.44502, 3.44995, 3.4768, 3.47568, 3.46933, 3.46542, 3.44973, 3.4116, 3.48635, 3.41002, 3.4635, 3.60598, 3.36312, 3.37688, 3.45509, 3.42031, 3.40811, 3.39346, 3.51968, 3.37968, 3.30681, 3.41998, 3.42048, 3.40245, 3.32106, 3.35407, 3.36684, 3.36996, 3.39425, 3.37352, 3.27134, 3.33416, 3.36043, 3.2468, 3.19623, 3.25601, 3.18945, 3.2042, 3.25372, 3.17041, 3.09408, 3.09549, 3.00257, 3.09268, 3.12798, 3.02715, 3.02147, 2.97063, 2.8702, 2.88779, 2.87573, 2.96007, 2.86396, 2.82503, 2.85046, 2.74288, 2.83987, 2.75592, 2.72686, 2.61559, 2.78691, 2.52969, 2.52376, 2.64935, 2.47175, 2.43027, 2.44512, 2.32001, 2.35346, 2.29702, 2.26902, 2.19594, 2.17455, 2.13723, 2.21828, 2.17045, 2.11348, 2.05186, 1.90909, 1.94368, 1.85471, 1.82315, 1.71818, 1.6914, 1.69966, 1.59071, 1.59465, 1.58102, 1.49242, 1.3664, 1.36725, 1.34164, 1.36005, 1.36834, 1.3094, 1.21212, 1.21973, 1.14436, 1.16586, 1.13201, 1.17614, 1.09794, 1.03729, 1.06728, 1.07631, 1.03906, 1.09235, 1.06696, 1.03104, 0.9723, 0.97222, 0.91632, 0.95428, 0.91088, 0.90223, 0.94564, 0.94636, 0.84496, 0.95977, 1.02479, 0.865, 0.86768, 0.84073, 0.80731, 0.91962, 0.72997, 0.85239, 0.75502, 0.67988, 0.75747, 0.68226, 0.69764, 0.68729, 0.6821, 0.5413, 0.57141, 0.46762, 0.45146, 0.30402], [3.7197, 3.75288, 3.73948, 3.77312, 3.81365, 3.78157, 3.66759, 3.72636, 3.74291, 3.74803, 3.73465, 3.70506, 3.83496, 3.72305, 3.72498, 3.69183, 3.75015, 3.8738, 3.70245, 3.7183, 3.7553, 3.78608, 3.64957, 3.79166, 3.81488, 3.67751, 3.81057, 3.73817, 3.75199, 3.69218, 3.80472, 3.73387, 3.7714, 3.73614, 3.80875, 3.61156, 3.72281, 3.63094, 3.72263, 3.82989, 3.72763, 3.64282, 3.72133, 3.65832, 3.73877, 3.75626, 3.69515, 3.6809, 3.69723, 3.66151, 3.66526, 3.66839, 3.82306, 3.70673, 3.65888, 3.7321, 3.69242, 3.72043, 3.64903, 3.65518, 3.65633, 3.54972, 3.68196, 3.59492, 3.65627, 3.55864, 3.69094, 3.66159, 3.60427, 3.68031, 3.61134, 3.59775, 3.69613, 3.59021, 3.53424, 3.57728, 3.51997, 3.60853, 3.54904, 3.6011, 3.57333, 3.52464, 3.60692, 3.57399, 3.49841, 3.54567, 3.56203, 3.59262, 3.45313, 3.57509, 3.49024, 3.53294, 3.47659, 3.52539, 3.4238, 3.31427, 3.36263, 3.33281, 3.45233, 3.47506, 3.33258, 3.35411, 3.28456, 3.25264, 3.32047, 3.25435, 3.23093, 3.2172, 3.16335, 3.16917, 3.02875, 3.05147, 3.12606, 3.14517, 3.0664, 3.0573, 2.93829, 2.95637, 2.94692, 2.89459, 2.88349, 3.00862, 2.86556, 2.82761, 2.85602, 2.74692, 2.57108, 2.72543, 2.58654, 2.51087, 2.47341, 2.46957, 2.44169, 2.38086, 2.33281, 2.42138, 2.37016, 2.34938, 2.22008, 2.12459, 2.21601, 2.1142, 1.96385, 1.97391, 2.01273, 1.95927, 1.84508, 1.88961, 1.77264, 1.81971, 1.62878, 1.63765, 1.63094, 1.61532, 1.61961, 1.56051, 1.54983, 1.44875, 1.41198, 1.42107, 1.46421, 1.45253, 1.27968, 1.23784, 1.2987, 1.28364, 1.34296, 1.2227, 1.26018, 1.20951, 1.24435, 1.22171, 1.18628, 1.19984, 1.15413, 1.21284, 1.14366, 1.1792, 1.10195, 1.0899, 1.1852, 1.07179, 1.09002, 1.13247, 1.00483, 1.05028, 0.98792, 1.08739, 0.93726, 0.98243, 1.05575, 0.93122, 0.95288, 0.83264, 0.82001, 0.68427, 0.69098, 0.56786, 0.55199, 0.42048]]} \ No newline at end of file diff --git a/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/utils.py b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a05eee433eac10792f8ab6bb75b7d3cf7027ed56 --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/utils.py @@ -0,0 +1,116 @@ +import torch +from typing import Optional +from datasets import load_dataset, Features, Sequence, Value + +def build_target_layer_ids(num_target_layers: int, num_draft_layers: int): + if num_draft_layers == 1: + return [(num_target_layers // 2)] + start = 1 + end = num_target_layers - 3 + span = end - start + target_layer_ids = [ + int(round(start + (i * span) / (num_draft_layers - 1))) + for i in range(num_draft_layers) + ] + return target_layer_ids + +def extract_context_feature( + hidden_states: list[torch.Tensor], + layer_ids: Optional[list[int]], +) -> torch.Tensor: + offset = 1 + selected_states = [] + for layer_id in layer_ids: + selected_states.append(hidden_states[layer_id + offset]) + target_hidden = torch.cat(selected_states, dim=-1) + return target_hidden + +def sample(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor: + if temperature < 1e-5: + return torch.argmax(logits, dim=-1) + bsz, seq_len, vocab_size = logits.shape + logits = logits.view(-1, vocab_size) + logits = logits / temperature + probs = torch.softmax(logits, dim=-1) + return torch.multinomial(probs, num_samples=1).view(bsz, seq_len) + +def load_and_process_dataset(data_name: str): + # Math datasets + if data_name == "gsm8k": + dataset = load_dataset("openai/gsm8k", "main", split="test") + prompt_fmt = "{question}\nPlease reason step by step, and put your final answer within \\boxed{{}}." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "math500": + dataset = load_dataset("HuggingFaceH4/MATH-500", split="test") + prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "aime24": + dataset = load_dataset("HuggingFaceH4/aime_2024", split="train") + prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "aime25": + dataset = load_dataset("MathArena/aime_2025", split="train") + prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + # Chat datasets + elif data_name == "alpaca": + dataset = load_dataset("tatsu-lab/alpaca", split="train") + dataset = dataset.map(lambda x: {"formatted_input": (f"{x['instruction']}\n\nInput:\n{x['input']}" if x['input'] else x['instruction'])}) + dataset = dataset.map(lambda x: {"turns": [x["formatted_input"]]}) + + elif data_name == "mt-bench": + dataset = load_dataset("HuggingFaceH4/mt_bench_prompts", split="train") + dataset = dataset.map(lambda x: {"turns": x["prompt"]}) + + # Coding datasets + elif data_name == "humaneval": + dataset = load_dataset("openai/openai_humaneval", split="test") + prompt_fmt = "Write a solution to the following problem and make sure that it passes the tests:\n```python\n{prompt}\n```" + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "mbpp": + dataset = load_dataset("google-research-datasets/mbpp", "sanitized", split="test") + dataset = dataset.map(lambda x: {"turns": [x["prompt"]]}) + + elif data_name == "lbpp": + LBPP_PY_TEST_URL = "https://huggingface.co/datasets/CohereLabs/lbpp/resolve/main/python/test.parquet" + dataset = load_dataset("parquet", data_files={"test": LBPP_PY_TEST_URL})["test"] + dataset = dataset.map(lambda x: {"turns": [x["instruction"]]}) + + elif data_name == "swe-bench": + dataset = load_dataset("princeton-nlp/SWE-bench_Lite", split="test") + prompt_fmt = "Problem Statement:\n{problem_statement}\nPlease fix the issue described above." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "livecodebench": + base = "https://huggingface.co/datasets/livecodebench/code_generation_lite/resolve/main/" + allowed_files = ["test.jsonl", "test2.jsonl", "test3.jsonl", "test4.jsonl", "test5.jsonl", "test6.jsonl"] + urls = [base + fn for fn in allowed_files] + dataset = load_dataset("json", data_files={"test": urls})["test"] + def format_lcb(doc): + system_prompt = ( + "You are an expert Python programmer. You will be given a question (problem specification) " + "and will generate a correct Python program that matches the specification and passes all tests. " + "You will NOT return anything except for the program" + ) + question_block = f"### Question:\n{doc['question_content']}" + if doc.get("starter_code"): + format_message = "### Format: Use the following code structure:" + code_block = f"```python\n{doc['starter_code']}\n```" + else: + format_message = "### Format: Write your code in the following format:" + code_block = "```python\n# YOUR CODE HERE\n```" + answer_footer = "### Answer: (use the provided format with backticks)" + return f"{system_prompt}\n\n{question_block}\n\n{format_message}\n{code_block}\n\n{answer_footer}" + target_features = Features({"turns": Sequence(Value("large_string"))}) + dataset = dataset.map( + lambda x: {"turns": [format_lcb(x)]}, + remove_columns=dataset.column_names, + features=target_features + ) + + return dataset \ No newline at end of file diff --git a/runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/flow.safetensors b/runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/flow.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..4e2135b81f268860856ef8561c2033ac1ab74139 --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/flow.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e6d49e24d000733a4b8d9ea5be4a4e660924163512515f45114af9c2551f2946 +size 1797091260 diff --git a/runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/head.safetensors b/runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/head.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e70752ec54d4f1d23ffab7c6b0ebc41e259e6f7c --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/head.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:39248be7fbbbf3833d300044811fe41539eaa3177644112b719816b519ec70f2 +size 311165128 diff --git a/runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/model.safetensors b/runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..ec3ded5a312eb5bbeeb49fdf0335e2cd695cc267 --- /dev/null +++ b/runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d843fb298102997a57d16a9d0924b5dfbb7058f30a5dc8045d39a3a9f8d72847 +size 1074860568 diff --git a/runs/fmlm2_v3_6ep_ce/export/epoch_1/dflash.py b/runs/fmlm2_v3_6ep_ce/export/epoch_1/dflash.py new file mode 100644 index 0000000000000000000000000000000000000000..f3cd5d26057e92f4721977cf53195c2d738a6b76 --- /dev/null +++ b/runs/fmlm2_v3_6ep_ce/export/epoch_1/dflash.py @@ -0,0 +1,277 @@ +from typing import Optional, Callable +from typing_extensions import Unpack, Tuple +import torch +from torch import nn +from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3RMSNorm, + Qwen3RotaryEmbedding, + Qwen3Config, + Qwen3PreTrainedModel, + Qwen3MLP, + GradientCheckpointingLayer, + FlashAttentionKwargs, + rotate_half, + eager_attention_forward, + ALL_ATTENTION_FUNCTIONS, +) +from transformers import DynamicCache +from transformers.modeling_outputs import CausalLMOutputWithPast +from transformers.cache_utils import Cache +from .utils import build_target_layer_ids, extract_context_feature, sample + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_len = q.size(-2) + q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :]) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + +class Qwen3DFlashAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = False + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None + + def forward( + self, + hidden_states: torch.Tensor, + target_hidden: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_values: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + bsz, q_len = hidden_states.shape[:-1] + ctx_len = target_hidden.shape[1] + q = self.q_proj(hidden_states) + q = q.view(bsz, q_len, -1, self.head_dim) + q = self.q_norm(q).transpose(1, 2) + k_ctx = self.k_proj(target_hidden) + k_noise = self.k_proj(hidden_states) + v_ctx = self.v_proj(target_hidden) + v_noise = self.v_proj(hidden_states) + k = torch.cat([k_ctx, k_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim) + v = torch.cat([v_ctx, v_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim) + k = self.k_norm(k).transpose(1, 2) + v = v.transpose(1, 2) + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb(q, k, cos, sin) + if past_key_values is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs) + attn_fn: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + attn_output, attn_weights = attn_fn( + self, + q, + k, + v, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, q_len, -1) + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + +class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Qwen3DFlashAttention(config=config, layer_idx=layer_idx) + self.mlp = Qwen3MLP(config) + self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + target_hidden: Optional[torch.Tensor] = None, + hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + hidden_states=hidden_states, + target_hidden=target_hidden, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + )[0] + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + +class DFlashDraftModel(Qwen3PreTrainedModel): + config_class = Qwen3Config + _no_split_modules = ["Qwen3DFlashDecoderLayer"] + + def __init__(self, config) -> None: + super().__init__(config) + self.config = config + self.layers = nn.ModuleList( + [Qwen3DFlashDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.target_layer_ids = self.config.dflash_config.get("target_layer_ids", build_target_layer_ids(config.num_target_layers, config.num_hidden_layers)) + self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3RotaryEmbedding(config) + self.fc = nn.Linear(len(self.target_layer_ids) * config.hidden_size, config.hidden_size, bias=False) + self.hidden_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.block_size = config.block_size + self.mask_token_id = self.config.dflash_config.get("mask_token_id", None) + self.post_init() + + def forward( + self, + position_ids: torch.LongTensor, + attention_mask: Optional[torch.Tensor] = None, + noise_embedding: Optional[torch.Tensor] = None, + target_hidden: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: bool = False, + **kwargs, + ) -> CausalLMOutputWithPast: + hidden_states = noise_embedding + target_hidden = self.hidden_norm(self.fc(target_hidden)) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for layer in self.layers: + hidden_states = layer( + hidden_states=hidden_states, + target_hidden=target_hidden, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + return self.norm(hidden_states) + + @torch.inference_mode() + def spec_generate( + self, + target: nn.Module, + input_ids: torch.LongTensor, + max_new_tokens: int, + stop_token_ids: list[int], + temperature: float, + ): + self.eval() + num_input_tokens = input_ids.shape[1] + max_length = num_input_tokens + max_new_tokens + + block_size = self.block_size + output_ids = torch.full( + (1, max_length + block_size), + self.mask_token_id, + dtype=torch.long, + device=target.device, + ) + position_ids = torch.arange(output_ids.shape[1], device=target.device).unsqueeze(0) + + past_key_values_target = DynamicCache() + past_key_values_draft = DynamicCache() + + # Prefill stage + output = target( + input_ids, + position_ids=position_ids[:, :num_input_tokens], + past_key_values=past_key_values_target, + use_cache=True, + logits_to_keep=1, + output_hidden_states=True, + ) + + output_ids[:, :num_input_tokens] = input_ids + output_ids[:, num_input_tokens:num_input_tokens+1] = sample(output.logits, temperature) + target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids) + + # Decode stage + acceptance_lengths = [] + start = input_ids.shape[1] + while start < max_length: + block_output_ids = output_ids[:, start : start + block_size].clone() + block_position_ids = position_ids[:, start : start + block_size] + noise_embedding = target.model.embed_tokens(block_output_ids) + draft_logits = target.lm_head(self( + target_hidden=target_hidden, + noise_embedding=noise_embedding, + position_ids=position_ids[:, past_key_values_draft.get_seq_length(): start + block_size], + past_key_values=past_key_values_draft, + use_cache=True, + is_causal=False, + )[:, -block_size+1:, :]) + past_key_values_draft.crop(start) + block_output_ids[:, 1:] = sample(draft_logits) + + output = target( + block_output_ids, + position_ids=block_position_ids, + past_key_values=past_key_values_target, + use_cache=True, + output_hidden_states=True, + ) + + posterior = sample(output.logits, temperature) + acceptance_length = (block_output_ids[:, 1:] == posterior[:, :-1]).cumprod(dim=1).sum(dim=1)[0].item() + output_ids[:, start : start + acceptance_length + 1] = block_output_ids[:, : acceptance_length + 1] + output_ids[:, start + acceptance_length + 1] = posterior[:, acceptance_length] + start += acceptance_length + 1 + past_key_values_target.crop(start) + target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids)[:, :acceptance_length + 1, :] + acceptance_lengths.append(acceptance_length+1) + if stop_token_ids is not None and any( + stop_token_id in output_ids[:, num_input_tokens:] for stop_token_id in stop_token_ids + ): + break + output_ids = output_ids[:, :max_length] + output_ids = output_ids[:, output_ids[0] != self.mask_token_id] + if stop_token_ids is not None: + stop_token_ids = torch.tensor(stop_token_ids, device=output_ids.device) + stop_token_indices = torch.isin(output_ids[0][num_input_tokens:], stop_token_ids).nonzero(as_tuple=True)[0] + if stop_token_indices.numel() > 0: + output_ids = output_ids[:, : num_input_tokens + stop_token_indices[0] + 1] + + return output_ids diff --git a/runs/fmlm2_v3_6ep_ce/export/epoch_1/flow.safetensors b/runs/fmlm2_v3_6ep_ce/export/epoch_1/flow.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..098dd68e07537c0c7e7633d6bfdb30cae4d4f668 --- /dev/null +++ b/runs/fmlm2_v3_6ep_ce/export/epoch_1/flow.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:418550eb1cc8e6b3e131beea42e9a7d52540fd63688c3b0da8cca9f8116c35c8 +size 1770876772 diff --git a/runs/fmlm2_v3_6ep_ce/export/epoch_1/flow_config.json b/runs/fmlm2_v3_6ep_ce/export/epoch_1/flow_config.json new file mode 100644 index 0000000000000000000000000000000000000000..229f80ae702ecf48464c0188f607da376b9216b2 --- /dev/null +++ b/runs/fmlm2_v3_6ep_ce/export/epoch_1/flow_config.json @@ -0,0 +1,10 @@ +{ + "learn_embed": true, + "epoch": 1, + "step": 2859, + "label": "epoch_1", + "objective": "fmlm2", + "kd": true, + "readout_frac": 0.25, + "head_rank": 256 +} \ No newline at end of file diff --git a/runs/fmlm2_v3_6ep_ce/export/epoch_1/head.safetensors b/runs/fmlm2_v3_6ep_ce/export/epoch_1/head.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e06c95de5e21074426a813873ddbb1c7b0663617 --- /dev/null +++ b/runs/fmlm2_v3_6ep_ce/export/epoch_1/head.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f5029aa2fcc79d191a60a735d3c368d6103fcd12e35d7268b6d3bf347fd8f8e3 +size 311165128 diff --git a/runs/fmlm2_v3_6ep_ce/export/epoch_1/model.safetensors b/runs/fmlm2_v3_6ep_ce/export/epoch_1/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..49bb41cb35c1bfb9fac7c9ed792cfb995ff58690 --- /dev/null +++ b/runs/fmlm2_v3_6ep_ce/export/epoch_1/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bb3cc23203f226b2b9df422557614d7119f684904244ef7c13b56f5d932f2e19 +size 1074860568 diff --git a/runs/fmlm2_v3_6ep_ce/export/epoch_1/utils.py b/runs/fmlm2_v3_6ep_ce/export/epoch_1/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a05eee433eac10792f8ab6bb75b7d3cf7027ed56 --- /dev/null +++ b/runs/fmlm2_v3_6ep_ce/export/epoch_1/utils.py @@ -0,0 +1,116 @@ +import torch +from typing import Optional +from datasets import load_dataset, Features, Sequence, Value + +def build_target_layer_ids(num_target_layers: int, num_draft_layers: int): + if num_draft_layers == 1: + return [(num_target_layers // 2)] + start = 1 + end = num_target_layers - 3 + span = end - start + target_layer_ids = [ + int(round(start + (i * span) / (num_draft_layers - 1))) + for i in range(num_draft_layers) + ] + return target_layer_ids + +def extract_context_feature( + hidden_states: list[torch.Tensor], + layer_ids: Optional[list[int]], +) -> torch.Tensor: + offset = 1 + selected_states = [] + for layer_id in layer_ids: + selected_states.append(hidden_states[layer_id + offset]) + target_hidden = torch.cat(selected_states, dim=-1) + return target_hidden + +def sample(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor: + if temperature < 1e-5: + return torch.argmax(logits, dim=-1) + bsz, seq_len, vocab_size = logits.shape + logits = logits.view(-1, vocab_size) + logits = logits / temperature + probs = torch.softmax(logits, dim=-1) + return torch.multinomial(probs, num_samples=1).view(bsz, seq_len) + +def load_and_process_dataset(data_name: str): + # Math datasets + if data_name == "gsm8k": + dataset = load_dataset("openai/gsm8k", "main", split="test") + prompt_fmt = "{question}\nPlease reason step by step, and put your final answer within \\boxed{{}}." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "math500": + dataset = load_dataset("HuggingFaceH4/MATH-500", split="test") + prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "aime24": + dataset = load_dataset("HuggingFaceH4/aime_2024", split="train") + prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "aime25": + dataset = load_dataset("MathArena/aime_2025", split="train") + prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + # Chat datasets + elif data_name == "alpaca": + dataset = load_dataset("tatsu-lab/alpaca", split="train") + dataset = dataset.map(lambda x: {"formatted_input": (f"{x['instruction']}\n\nInput:\n{x['input']}" if x['input'] else x['instruction'])}) + dataset = dataset.map(lambda x: {"turns": [x["formatted_input"]]}) + + elif data_name == "mt-bench": + dataset = load_dataset("HuggingFaceH4/mt_bench_prompts", split="train") + dataset = dataset.map(lambda x: {"turns": x["prompt"]}) + + # Coding datasets + elif data_name == "humaneval": + dataset = load_dataset("openai/openai_humaneval", split="test") + prompt_fmt = "Write a solution to the following problem and make sure that it passes the tests:\n```python\n{prompt}\n```" + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "mbpp": + dataset = load_dataset("google-research-datasets/mbpp", "sanitized", split="test") + dataset = dataset.map(lambda x: {"turns": [x["prompt"]]}) + + elif data_name == "lbpp": + LBPP_PY_TEST_URL = "https://huggingface.co/datasets/CohereLabs/lbpp/resolve/main/python/test.parquet" + dataset = load_dataset("parquet", data_files={"test": LBPP_PY_TEST_URL})["test"] + dataset = dataset.map(lambda x: {"turns": [x["instruction"]]}) + + elif data_name == "swe-bench": + dataset = load_dataset("princeton-nlp/SWE-bench_Lite", split="test") + prompt_fmt = "Problem Statement:\n{problem_statement}\nPlease fix the issue described above." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "livecodebench": + base = "https://huggingface.co/datasets/livecodebench/code_generation_lite/resolve/main/" + allowed_files = ["test.jsonl", "test2.jsonl", "test3.jsonl", "test4.jsonl", "test5.jsonl", "test6.jsonl"] + urls = [base + fn for fn in allowed_files] + dataset = load_dataset("json", data_files={"test": urls})["test"] + def format_lcb(doc): + system_prompt = ( + "You are an expert Python programmer. You will be given a question (problem specification) " + "and will generate a correct Python program that matches the specification and passes all tests. " + "You will NOT return anything except for the program" + ) + question_block = f"### Question:\n{doc['question_content']}" + if doc.get("starter_code"): + format_message = "### Format: Use the following code structure:" + code_block = f"```python\n{doc['starter_code']}\n```" + else: + format_message = "### Format: Write your code in the following format:" + code_block = "```python\n# YOUR CODE HERE\n```" + answer_footer = "### Answer: (use the provided format with backticks)" + return f"{system_prompt}\n\n{question_block}\n\n{format_message}\n{code_block}\n\n{answer_footer}" + target_features = Features({"turns": Sequence(Value("large_string"))}) + dataset = dataset.map( + lambda x: {"turns": [format_lcb(x)]}, + remove_columns=dataset.column_names, + features=target_features + ) + + return dataset \ No newline at end of file diff --git a/runs/fmlm2_v3_6ep_tv/export/epoch_1/config.json b/runs/fmlm2_v3_6ep_tv/export/epoch_1/config.json new file mode 100644 index 0000000000000000000000000000000000000000..89073407b95c2f29693c309b6f00189c99b32b61 --- /dev/null +++ b/runs/fmlm2_v3_6ep_tv/export/epoch_1/config.json @@ -0,0 +1,52 @@ +{ + "architectures": [ + "DFlashDraftModel" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": { + "AutoModel": "dflash.DFlashDraftModel" + }, + "block_size": 16, + "bos_token_id": 151643, + "dflash_config": { + "mask_token_id": 151669, + "target_layer_ids": [ + 1, + 9, + 17, + 25, + 33 + ] + }, + "dtype": "bfloat16", + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 9728, + "layer_types": [ + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention" + ], + "max_position_embeddings": 40960, + "max_window_layers": 5, + "model_type": "qwen3", + "num_attention_heads": 32, + "num_hidden_layers": 5, + "num_key_value_heads": 8, + "num_target_layers": 36, + "rms_norm_eps": 1e-06, + "rope_scaling": null, + "rope_theta": 1000000, + "sliding_window": null, + "tie_word_embeddings": true, + "transformers_version": "4.57.3", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936 +} diff --git a/runs/fmlm2_v3_6ep_tv/export/epoch_1/dflash.py b/runs/fmlm2_v3_6ep_tv/export/epoch_1/dflash.py new file mode 100644 index 0000000000000000000000000000000000000000..f3cd5d26057e92f4721977cf53195c2d738a6b76 --- /dev/null +++ b/runs/fmlm2_v3_6ep_tv/export/epoch_1/dflash.py @@ -0,0 +1,277 @@ +from typing import Optional, Callable +from typing_extensions import Unpack, Tuple +import torch +from torch import nn +from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3RMSNorm, + Qwen3RotaryEmbedding, + Qwen3Config, + Qwen3PreTrainedModel, + Qwen3MLP, + GradientCheckpointingLayer, + FlashAttentionKwargs, + rotate_half, + eager_attention_forward, + ALL_ATTENTION_FUNCTIONS, +) +from transformers import DynamicCache +from transformers.modeling_outputs import CausalLMOutputWithPast +from transformers.cache_utils import Cache +from .utils import build_target_layer_ids, extract_context_feature, sample + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_len = q.size(-2) + q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :]) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + +class Qwen3DFlashAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = False + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None + + def forward( + self, + hidden_states: torch.Tensor, + target_hidden: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_values: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + bsz, q_len = hidden_states.shape[:-1] + ctx_len = target_hidden.shape[1] + q = self.q_proj(hidden_states) + q = q.view(bsz, q_len, -1, self.head_dim) + q = self.q_norm(q).transpose(1, 2) + k_ctx = self.k_proj(target_hidden) + k_noise = self.k_proj(hidden_states) + v_ctx = self.v_proj(target_hidden) + v_noise = self.v_proj(hidden_states) + k = torch.cat([k_ctx, k_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim) + v = torch.cat([v_ctx, v_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim) + k = self.k_norm(k).transpose(1, 2) + v = v.transpose(1, 2) + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb(q, k, cos, sin) + if past_key_values is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs) + attn_fn: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + attn_output, attn_weights = attn_fn( + self, + q, + k, + v, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, q_len, -1) + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + +class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Qwen3DFlashAttention(config=config, layer_idx=layer_idx) + self.mlp = Qwen3MLP(config) + self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + target_hidden: Optional[torch.Tensor] = None, + hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + hidden_states=hidden_states, + target_hidden=target_hidden, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + )[0] + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + +class DFlashDraftModel(Qwen3PreTrainedModel): + config_class = Qwen3Config + _no_split_modules = ["Qwen3DFlashDecoderLayer"] + + def __init__(self, config) -> None: + super().__init__(config) + self.config = config + self.layers = nn.ModuleList( + [Qwen3DFlashDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.target_layer_ids = self.config.dflash_config.get("target_layer_ids", build_target_layer_ids(config.num_target_layers, config.num_hidden_layers)) + self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3RotaryEmbedding(config) + self.fc = nn.Linear(len(self.target_layer_ids) * config.hidden_size, config.hidden_size, bias=False) + self.hidden_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.block_size = config.block_size + self.mask_token_id = self.config.dflash_config.get("mask_token_id", None) + self.post_init() + + def forward( + self, + position_ids: torch.LongTensor, + attention_mask: Optional[torch.Tensor] = None, + noise_embedding: Optional[torch.Tensor] = None, + target_hidden: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: bool = False, + **kwargs, + ) -> CausalLMOutputWithPast: + hidden_states = noise_embedding + target_hidden = self.hidden_norm(self.fc(target_hidden)) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for layer in self.layers: + hidden_states = layer( + hidden_states=hidden_states, + target_hidden=target_hidden, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + return self.norm(hidden_states) + + @torch.inference_mode() + def spec_generate( + self, + target: nn.Module, + input_ids: torch.LongTensor, + max_new_tokens: int, + stop_token_ids: list[int], + temperature: float, + ): + self.eval() + num_input_tokens = input_ids.shape[1] + max_length = num_input_tokens + max_new_tokens + + block_size = self.block_size + output_ids = torch.full( + (1, max_length + block_size), + self.mask_token_id, + dtype=torch.long, + device=target.device, + ) + position_ids = torch.arange(output_ids.shape[1], device=target.device).unsqueeze(0) + + past_key_values_target = DynamicCache() + past_key_values_draft = DynamicCache() + + # Prefill stage + output = target( + input_ids, + position_ids=position_ids[:, :num_input_tokens], + past_key_values=past_key_values_target, + use_cache=True, + logits_to_keep=1, + output_hidden_states=True, + ) + + output_ids[:, :num_input_tokens] = input_ids + output_ids[:, num_input_tokens:num_input_tokens+1] = sample(output.logits, temperature) + target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids) + + # Decode stage + acceptance_lengths = [] + start = input_ids.shape[1] + while start < max_length: + block_output_ids = output_ids[:, start : start + block_size].clone() + block_position_ids = position_ids[:, start : start + block_size] + noise_embedding = target.model.embed_tokens(block_output_ids) + draft_logits = target.lm_head(self( + target_hidden=target_hidden, + noise_embedding=noise_embedding, + position_ids=position_ids[:, past_key_values_draft.get_seq_length(): start + block_size], + past_key_values=past_key_values_draft, + use_cache=True, + is_causal=False, + )[:, -block_size+1:, :]) + past_key_values_draft.crop(start) + block_output_ids[:, 1:] = sample(draft_logits) + + output = target( + block_output_ids, + position_ids=block_position_ids, + past_key_values=past_key_values_target, + use_cache=True, + output_hidden_states=True, + ) + + posterior = sample(output.logits, temperature) + acceptance_length = (block_output_ids[:, 1:] == posterior[:, :-1]).cumprod(dim=1).sum(dim=1)[0].item() + output_ids[:, start : start + acceptance_length + 1] = block_output_ids[:, : acceptance_length + 1] + output_ids[:, start + acceptance_length + 1] = posterior[:, acceptance_length] + start += acceptance_length + 1 + past_key_values_target.crop(start) + target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids)[:, :acceptance_length + 1, :] + acceptance_lengths.append(acceptance_length+1) + if stop_token_ids is not None and any( + stop_token_id in output_ids[:, num_input_tokens:] for stop_token_id in stop_token_ids + ): + break + output_ids = output_ids[:, :max_length] + output_ids = output_ids[:, output_ids[0] != self.mask_token_id] + if stop_token_ids is not None: + stop_token_ids = torch.tensor(stop_token_ids, device=output_ids.device) + stop_token_indices = torch.isin(output_ids[0][num_input_tokens:], stop_token_ids).nonzero(as_tuple=True)[0] + if stop_token_indices.numel() > 0: + output_ids = output_ids[:, : num_input_tokens + stop_token_indices[0] + 1] + + return output_ids diff --git a/runs/fmlm2_v3_6ep_tv/export/epoch_1/dspark_config.json b/runs/fmlm2_v3_6ep_tv/export/epoch_1/dspark_config.json new file mode 100644 index 0000000000000000000000000000000000000000..f3a601e825cc50e7c627234fd1a2c56d7ec543b3 --- /dev/null +++ b/runs/fmlm2_v3_6ep_tv/export/epoch_1/dspark_config.json @@ -0,0 +1,16 @@ +{ + "rank": 256, + "backbone": { + "kind": "flow", + "path": "runs/fmlm2_v3_6ep_tv/export/epoch_1", + "mode": "onejump_zero" + }, + "epoch": 1, + "step": 2859, + "label": "epoch_1", + "joint": true, + "objective": "fmlm2", + "kd": true, + "readout_frac": 0.25, + "head_rank": 256 +} \ No newline at end of file diff --git a/runs/fmlm2_v3_6ep_tv/export/epoch_1/flow.safetensors b/runs/fmlm2_v3_6ep_tv/export/epoch_1/flow.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0469de950d6512c96e03a3c31a5d7f7911b197dc --- /dev/null +++ b/runs/fmlm2_v3_6ep_tv/export/epoch_1/flow.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:44c4455cbf38f0838eb86bc64d2cd1f92a41ebde5eb813e4b3aac298abc10ede +size 1770876772 diff --git a/runs/fmlm2_v3_6ep_tv/export/epoch_1/flow_config.json b/runs/fmlm2_v3_6ep_tv/export/epoch_1/flow_config.json new file mode 100644 index 0000000000000000000000000000000000000000..229f80ae702ecf48464c0188f607da376b9216b2 --- /dev/null +++ b/runs/fmlm2_v3_6ep_tv/export/epoch_1/flow_config.json @@ -0,0 +1,10 @@ +{ + "learn_embed": true, + "epoch": 1, + "step": 2859, + "label": "epoch_1", + "objective": "fmlm2", + "kd": true, + "readout_frac": 0.25, + "head_rank": 256 +} \ No newline at end of file diff --git a/runs/fmlm2_v3_6ep_tv/export/epoch_1/head.safetensors b/runs/fmlm2_v3_6ep_tv/export/epoch_1/head.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..52a394da92e7af33e0b6384f18071c3275a22e02 --- /dev/null +++ b/runs/fmlm2_v3_6ep_tv/export/epoch_1/head.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b2a33a09f0e5a5804b0e154d8dab1ee9b25a1d8273bd27609495ce3f8f7ee906 +size 311165128 diff --git a/runs/fmlm2_v3_6ep_tv/export/epoch_1/model.safetensors b/runs/fmlm2_v3_6ep_tv/export/epoch_1/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..fd5b9b1e741e9b7ba15cc9bcb5485f81309c1b6a --- /dev/null +++ b/runs/fmlm2_v3_6ep_tv/export/epoch_1/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dba373f994baac3322c5f8f5b771d6aceb2edff9551ddcaa0bbbe40665755f9f +size 1074860568 diff --git a/runs/fmlm2_v3_6ep_tv/export/epoch_1/modeling_dflash.py b/runs/fmlm2_v3_6ep_tv/export/epoch_1/modeling_dflash.py new file mode 100644 index 0000000000000000000000000000000000000000..ec9f19ffd11e242b0dd9e9469817f21ef31c2b04 --- /dev/null +++ b/runs/fmlm2_v3_6ep_tv/export/epoch_1/modeling_dflash.py @@ -0,0 +1,309 @@ +from typing import Optional, Callable +from typing_extensions import Unpack, Tuple +import torch +from torch import nn +from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3RMSNorm, + Qwen3RotaryEmbedding, + Qwen3Config, + Qwen3PreTrainedModel, + Qwen3MLP, + GradientCheckpointingLayer, + FlashAttentionKwargs, + rotate_half, + eager_attention_forward, + ALL_ATTENTION_FUNCTIONS, +) +from transformers import DynamicCache +from transformers.modeling_outputs import CausalLMOutputWithPast +from transformers.cache_utils import Cache + +def sample(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor: + if temperature < 1e-5: + return torch.argmax(logits, dim=-1) + bsz, seq_len, vocab_size = logits.shape + logits = logits.view(-1, vocab_size) + logits = logits / temperature + probs = torch.softmax(logits, dim=-1) + return torch.multinomial(probs, num_samples=1).view(bsz, seq_len) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_len = q.size(-2) + q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :]) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + +class Qwen3DFlashAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = False + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None + + def forward( + self, + hidden_states: torch.Tensor, + target_hidden: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_values: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + bsz, q_len = hidden_states.shape[:-1] + ctx_len = target_hidden.shape[1] + q = self.q_proj(hidden_states) + q = q.view(bsz, q_len, -1, self.head_dim) + q = self.q_norm(q).transpose(1, 2) + k_ctx = self.k_proj(target_hidden) + k_noise = self.k_proj(hidden_states) + v_ctx = self.v_proj(target_hidden) + v_noise = self.v_proj(hidden_states) + k = torch.cat([k_ctx, k_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim) + v = torch.cat([v_ctx, v_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim) + k = self.k_norm(k).transpose(1, 2) + v = v.transpose(1, 2) + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb(q, k, cos, sin) + if past_key_values is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs) + attn_fn: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + attn_output, attn_weights = attn_fn( + self, + q, + k, + v, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, q_len, -1) + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + +class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Qwen3DFlashAttention(config=config, layer_idx=layer_idx) + self.mlp = Qwen3MLP(config) + self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + target_hidden: Optional[torch.Tensor] = None, + hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + hidden_states=hidden_states, + target_hidden=target_hidden, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + )[0] + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + +def build_target_layer_ids(num_target_layers: int, num_draft_layers: int): + if num_draft_layers == 1: + return [(num_target_layers // 2)] + start = 1 + end = num_target_layers - 3 + span = end - start + target_layer_ids = [ + int(round(start + (i * span) / (num_draft_layers - 1))) + for i in range(num_draft_layers) + ] + return target_layer_ids + +def extract_context_feature( + hidden_states: list[torch.Tensor], + layer_ids: Optional[list[int]], +) -> torch.Tensor: + offset = 1 + selected_states = [] + for layer_id in layer_ids: + selected_states.append(hidden_states[layer_id + offset]) + target_hidden = torch.cat(selected_states, dim=-1) + return target_hidden + +class DFlashDraftModel(Qwen3PreTrainedModel): + config_class = Qwen3Config + _no_split_modules = ["Qwen3DFlashDecoderLayer"] + def __init__(self, config) -> None: + super().__init__(config) + self.config = config + self.layers = nn.ModuleList( + [Qwen3DFlashDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.target_layer_ids = build_target_layer_ids(config.num_target_layers, config.num_hidden_layers) + self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3RotaryEmbedding(config) + self.fc = nn.Linear(len(self.target_layer_ids) * config.hidden_size, config.hidden_size, bias=False) + self.hidden_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.block_size = config.block_size + self.post_init() + + def forward( + self, + position_ids: torch.LongTensor, + attention_mask: Optional[torch.Tensor] = None, + noise_embedding: Optional[torch.Tensor] = None, + target_hidden: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: bool = False, + **kwargs, + ) -> CausalLMOutputWithPast: + hidden_states = noise_embedding + target_hidden = self.hidden_norm(self.fc(target_hidden)) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for layer in self.layers: + hidden_states = layer( + hidden_states=hidden_states, + target_hidden=target_hidden, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + return self.norm(hidden_states) + + @torch.inference_mode() + def spec_generate( + self, + target: nn.Module, + input_ids: torch.LongTensor, + mask_token_id: int, + max_new_tokens: int, + stop_token_ids: list[int], + temperature: float, + ): + self.eval() + target.eval() + num_input_tokens = input_ids.shape[1] + max_length = num_input_tokens + max_new_tokens + + block_size = self.block_size + output_ids = torch.full( + (1, max_length + block_size), + mask_token_id, + dtype=torch.long, + device=target.device, + ) + position_ids = torch.arange(output_ids.shape[1], device=target.device).unsqueeze(0) + + past_key_values_target = DynamicCache() + past_key_values_draft = DynamicCache() + + # Prefill stage + output = target( + input_ids, + position_ids=position_ids[:, :num_input_tokens], + past_key_values=past_key_values_target, + use_cache=True, + logits_to_keep=1, + output_hidden_states=True, + ) + + output_ids[:, :num_input_tokens] = input_ids + output_ids[:, num_input_tokens:num_input_tokens+1] = sample(output.logits, temperature) + target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids) + + # Decode stage + acceptance_lengths = [] + start = input_ids.shape[1] + while start < max_length: + block_output_ids = output_ids[:, start : start + block_size].clone() + block_position_ids = position_ids[:, start : start + block_size] + noise_embedding = target.model.embed_tokens(block_output_ids) + draft_logits = target.lm_head(self( + target_hidden=target_hidden, + noise_embedding=noise_embedding, + position_ids=position_ids[:, past_key_values_draft.get_seq_length(): start + block_size], + past_key_values=past_key_values_draft, + use_cache=True, + is_causal=False, + )[:, -block_size+1:, :]) + past_key_values_draft.crop(start) + block_output_ids[:, 1:] = sample(draft_logits) + + output = target( + block_output_ids, + position_ids=block_position_ids, + past_key_values=past_key_values_target, + use_cache=True, + output_hidden_states=True, + ) + + posterior = sample(output.logits, temperature) + acceptance_length = (block_output_ids[:, 1:] == posterior[:, :-1]).cumprod(dim=1).sum(dim=1)[0].item() + output_ids[:, start : start + acceptance_length + 1] = block_output_ids[:, : acceptance_length + 1] + output_ids[:, start + acceptance_length + 1] = posterior[:, acceptance_length] + start += acceptance_length + 1 + past_key_values_target.crop(start) + target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids)[:, :acceptance_length + 1, :] + acceptance_lengths.append(acceptance_length+1) + if stop_token_ids is not None and any( + stop_token_id in output_ids[:, num_input_tokens:] for stop_token_id in stop_token_ids + ): + break + output_ids = output_ids[:, :max_length] + output_ids = output_ids[:, output_ids[0] != mask_token_id] + if stop_token_ids is not None: + stop_token_ids = torch.tensor(stop_token_ids, device=output_ids.device) + stop_token_indices = torch.isin(output_ids[0][num_input_tokens:], stop_token_ids).nonzero(as_tuple=True)[0] + if stop_token_indices.numel() > 0: + output_ids = output_ids[:, : num_input_tokens + stop_token_indices[0] + 1] + + return output_ids diff --git a/runs/fmlm2_v3_6ep_tv/export/epoch_1/utils.py b/runs/fmlm2_v3_6ep_tv/export/epoch_1/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a05eee433eac10792f8ab6bb75b7d3cf7027ed56 --- /dev/null +++ b/runs/fmlm2_v3_6ep_tv/export/epoch_1/utils.py @@ -0,0 +1,116 @@ +import torch +from typing import Optional +from datasets import load_dataset, Features, Sequence, Value + +def build_target_layer_ids(num_target_layers: int, num_draft_layers: int): + if num_draft_layers == 1: + return [(num_target_layers // 2)] + start = 1 + end = num_target_layers - 3 + span = end - start + target_layer_ids = [ + int(round(start + (i * span) / (num_draft_layers - 1))) + for i in range(num_draft_layers) + ] + return target_layer_ids + +def extract_context_feature( + hidden_states: list[torch.Tensor], + layer_ids: Optional[list[int]], +) -> torch.Tensor: + offset = 1 + selected_states = [] + for layer_id in layer_ids: + selected_states.append(hidden_states[layer_id + offset]) + target_hidden = torch.cat(selected_states, dim=-1) + return target_hidden + +def sample(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor: + if temperature < 1e-5: + return torch.argmax(logits, dim=-1) + bsz, seq_len, vocab_size = logits.shape + logits = logits.view(-1, vocab_size) + logits = logits / temperature + probs = torch.softmax(logits, dim=-1) + return torch.multinomial(probs, num_samples=1).view(bsz, seq_len) + +def load_and_process_dataset(data_name: str): + # Math datasets + if data_name == "gsm8k": + dataset = load_dataset("openai/gsm8k", "main", split="test") + prompt_fmt = "{question}\nPlease reason step by step, and put your final answer within \\boxed{{}}." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "math500": + dataset = load_dataset("HuggingFaceH4/MATH-500", split="test") + prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "aime24": + dataset = load_dataset("HuggingFaceH4/aime_2024", split="train") + prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "aime25": + dataset = load_dataset("MathArena/aime_2025", split="train") + prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + # Chat datasets + elif data_name == "alpaca": + dataset = load_dataset("tatsu-lab/alpaca", split="train") + dataset = dataset.map(lambda x: {"formatted_input": (f"{x['instruction']}\n\nInput:\n{x['input']}" if x['input'] else x['instruction'])}) + dataset = dataset.map(lambda x: {"turns": [x["formatted_input"]]}) + + elif data_name == "mt-bench": + dataset = load_dataset("HuggingFaceH4/mt_bench_prompts", split="train") + dataset = dataset.map(lambda x: {"turns": x["prompt"]}) + + # Coding datasets + elif data_name == "humaneval": + dataset = load_dataset("openai/openai_humaneval", split="test") + prompt_fmt = "Write a solution to the following problem and make sure that it passes the tests:\n```python\n{prompt}\n```" + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "mbpp": + dataset = load_dataset("google-research-datasets/mbpp", "sanitized", split="test") + dataset = dataset.map(lambda x: {"turns": [x["prompt"]]}) + + elif data_name == "lbpp": + LBPP_PY_TEST_URL = "https://huggingface.co/datasets/CohereLabs/lbpp/resolve/main/python/test.parquet" + dataset = load_dataset("parquet", data_files={"test": LBPP_PY_TEST_URL})["test"] + dataset = dataset.map(lambda x: {"turns": [x["instruction"]]}) + + elif data_name == "swe-bench": + dataset = load_dataset("princeton-nlp/SWE-bench_Lite", split="test") + prompt_fmt = "Problem Statement:\n{problem_statement}\nPlease fix the issue described above." + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) + + elif data_name == "livecodebench": + base = "https://huggingface.co/datasets/livecodebench/code_generation_lite/resolve/main/" + allowed_files = ["test.jsonl", "test2.jsonl", "test3.jsonl", "test4.jsonl", "test5.jsonl", "test6.jsonl"] + urls = [base + fn for fn in allowed_files] + dataset = load_dataset("json", data_files={"test": urls})["test"] + def format_lcb(doc): + system_prompt = ( + "You are an expert Python programmer. You will be given a question (problem specification) " + "and will generate a correct Python program that matches the specification and passes all tests. " + "You will NOT return anything except for the program" + ) + question_block = f"### Question:\n{doc['question_content']}" + if doc.get("starter_code"): + format_message = "### Format: Use the following code structure:" + code_block = f"```python\n{doc['starter_code']}\n```" + else: + format_message = "### Format: Write your code in the following format:" + code_block = "```python\n# YOUR CODE HERE\n```" + answer_footer = "### Answer: (use the provided format with backticks)" + return f"{system_prompt}\n\n{question_block}\n\n{format_message}\n{code_block}\n\n{answer_footer}" + target_features = Features({"turns": Sequence(Value("large_string"))}) + dataset = dataset.map( + lambda x: {"turns": [format_lcb(x)]}, + remove_columns=dataset.column_names, + features=target_features + ) + + return dataset \ No newline at end of file diff --git a/runs/fmlm2_v3_tv1p/ckpt/step_1250.pt b/runs/fmlm2_v3_tv1p/ckpt/step_1250.pt new file mode 100644 index 0000000000000000000000000000000000000000..8b19f96fda2c14d5c93f78989a6a641aeccd046d --- /dev/null +++ b/runs/fmlm2_v3_tv1p/ckpt/step_1250.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9071b4ab56db18de4722350a15f045102032589e7c1793030b517db7f6fd5dd1 +size 12695352064 diff --git a/runs/full_scratch/export/epoch_6/flow.safetensors b/runs/full_scratch/export/epoch_6/flow.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..32e0fe83fdf7ecf6748f423c6bb1d1475b0237eb --- /dev/null +++ b/runs/full_scratch/export/epoch_6/flow.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d9e232bedcd5475f9de6bb35fe7ef49f6fb04bf487bc8d1861a56b84a9ef3d30 +size 1770876772 diff --git a/runs/full_scratch/export/epoch_6/model.safetensors b/runs/full_scratch/export/epoch_6/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..6ff0f87d58570bd4cebd7ea08918f7dedbf59d05 --- /dev/null +++ b/runs/full_scratch/export/epoch_6/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:698112c5e83c2e4bcad11b0d3328cba7576cc146c332e6fb8d58e1ff31f0a9c7 +size 1074860568