add in example and run pre-commit

Author: wpybtw

README.md

@@ -164,6 +164,9 @@ bash ./examples/run_qwen3_moe_eagle3_online.sh
 
 # train Qwen3-8B
 bash ./examples/run_qwen3_dense_eagle3_online.sh
+
+# train Qwq-32B
+bash ./examples/run_qwq_dense_eagle3_online.sh
 ```
 
 ### 💨 Offline Training

scripts/train_eagle3_offline.py

@@ -4,14 +4,14 @@
 
 import torch
 import torch.distributed as dist
-import wandb
 from accelerate.utils import set_seed
 from datasets import load_dataset
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp import MixedPrecision, ShardingStrategy, StateDictType
 from tqdm import tqdm
 from transformers import AutoTokenizer
 
+import wandb
 from specforge import AutoDraftModelConfig, AutoEagle3DraftModel, OfflineEagle3Model
 from specforge.data import (
     build_eagle3_dataset,

scripts/train_eagle3_online.py

@@ -4,14 +4,14 @@
 
 import torch
 import torch.distributed as dist
-import wandb
 from accelerate.utils import set_seed
 from datasets import load_dataset
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp import MixedPrecision, ShardingStrategy, StateDictType
 from tqdm import tqdm
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+import wandb
 from specforge import (
     AutoDistributedTargetModel,
     AutoDraftModelConfig,

specforge/modeling/auto.py

@@ -10,18 +10,17 @@
     Llama4TextConfig,
     LlamaConfig,
     PretrainedConfig,
-    Qwen3MoeConfig,
     Qwen2Config,
+    Qwen3MoeConfig,
 )
 
 from specforge.utils import default_torch_dtype
 
 from .draft.llama3_eagle import LlamaForCausalLMEagle3
-from .target.llama4 import Llama4ForCausalLM
-from .target.qwen3_moe import Qwen3MoeForCausalLM
-
 from .draft.qwen2_eagle import Qwen2ForCausalLMEagle3
+from .target.llama4 import Llama4ForCausalLM
 from .target.qwen2 import Qwen2ForCausalLM
+from .target.qwen3_moe import Qwen3MoeForCausalLM
 
 
 class AutoEagle3DraftModel(AutoModelForCausalLMBase):

specforge/modeling/draft/qwen2_eagle.py

@@ -5,7 +5,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import GenerationMixin, Qwen2Config, PreTrainedModel
+from transformers import GenerationMixin, PreTrainedModel, Qwen2Config
 from transformers.activations import ACT2FN
 from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
 

specforge/modeling/target/qwen2.py

@@ -22,7 +22,10 @@
 from transformers.cache_utils import Cache, DynamicCache
 from transformers.generation import GenerationMixin
 from transformers.integrations import use_kernel_forward_from_hub
-from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.masking_utils import (
+    create_causal_mask,
+    create_sliding_window_causal_mask,
+)
 from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
 from transformers.modeling_layers import GradientCheckpointingLayer
 from transformers.modeling_outputs import (
@@ -146,7 +149,9 @@ def forward(
         value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
 
         cos, sin = position_embeddings
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
 
         if past_key_value is not None:
             # sin and cos are specific to RoPE models; cache_position needed for the static cache
@@ -187,9 +192,7 @@ def __init__(self, config: Qwen2Config, layer_idx: int):
         self.self_attn = Qwen2Attention(config=config, layer_idx=layer_idx)
 
         self.mlp = Qwen2MLP(config)
-        self.input_layernorm = Qwen2RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps
-        )
+        self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = Qwen2RMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
         )
@@ -381,9 +384,9 @@ def forward(
             }
             # The sliding window alternating layers are not always activated depending on the config
             if self.has_sliding_layers:
-                causal_mask_mapping[
-                    "sliding_attention"
-                ] = create_sliding_window_causal_mask(**mask_kwargs)
+                causal_mask_mapping["sliding_attention"] = (
+                    create_sliding_window_causal_mask(**mask_kwargs)
+                )
 
         hidden_states = inputs_embeds
 
@@ -843,4 +846,4 @@ def forward(
     "Qwen2ForSequenceClassification",
     "Qwen2ForTokenClassification",
     "Qwen2ForQuestionAnswering",
-]
+]

tests/test_qwen2_tp.py

@@ -6,7 +6,7 @@
 import torch.distributed as dist
 import torch.multiprocessing as mp
 from accelerate.utils import set_seed
-from transformers import Qwen2ForCausalLM, Qwen2Config
+from transformers import Qwen2Config, Qwen2ForCausalLM
 
 from specforge.distributed import init_distributed
 
@@ -38,9 +38,7 @@ def test_qwen2_tp(rank, world_size, temp_dir):
     # create the single-gpu
     model = Qwen2ForCausalLM(config).cuda()
 
-    from specforge.modeling.target.qwen2 import (
-        Qwen2ForCausalLM as DistQwen2ForCausalLM,
-    )
+    from specforge.modeling.target.qwen2 import Qwen2ForCausalLM as DistQwen2ForCausalLM
 
     dist_model = DistQwen2ForCausalLM(config).cuda()