fix bugs

Author: tastelikefeet

cookbook/megatron/tp.py

@@ -30,7 +30,7 @@ def eval(model):
     dataset.set_template('Template', model_id='ms://Qwen/Qwen2.5-7B-Instruct')
     dataset.map(SelfCognitionProcessor('twinkle大模型', 'ModelScope社区'))
     dataset.encode()
-    dataloader = DataLoader(dataset=dataset, batch_size=1)
+    dataloader = DataLoader(dataset=dataset, batch_size=16)
     for step, batch in tqdm(enumerate(dataloader)):
         model.forward_only(inputs=batch)
     metrics = model.calculate_metric(is_training=False)

cookbook/megatron/tp_moe.py

@@ -30,7 +30,7 @@ def eval(model):
     dataset.set_template('Template', model_id='ms://Qwen/Qwen3-30B-A3B-Instruct-2507')
     dataset.map(SelfCognitionProcessor('twinkle大模型', 'ModelScope社区'))
     dataset.encode()
-    dataloader = DataLoader(dataset=dataset, batch_size=1)
+    dataloader = DataLoader(dataset=dataset, batch_size=16)
     for step, batch in tqdm(enumerate(dataloader)):
         model.forward_only(inputs=batch)
     metrics = model.calculate_metric(is_training=False)

cookbook/rl/grpo.py

@@ -37,8 +37,8 @@
 ADAPTER_NAME = 'default'
 
 def create_gsm8k_dataset():
-    dataset = Dataset(DatasetMeta("ms://modelscope/gsm8k", subset_name='main', split='train'))
-    dataset.set_template("Template", model_id=MODEL_ID, max_length=2048)
+    dataset = Dataset(DatasetMeta('ms://modelscope/gsm8k', subset_name='main', split='train'))
+    dataset.set_template('Template', model_id=MODEL_ID, max_length=2048)
     dataset.map(GSM8KProcessor())
     dataset.encode(add_generation_prompt=True)
     return dataset
@@ -67,7 +67,7 @@ def main():
     sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS)
     twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, lazy_collect=False)
 
-    lora_config = LoraConfig(target_modules="all-linear", r=32, lora_alpha=64, lora_dropout=0.05)
+    lora_config = LoraConfig(target_modules='all-linear', r=32, lora_alpha=64, lora_dropout=0.05)
 
     if USE_MEGATRON:
         from twinkle.model.megatron import MegatronModel
@@ -164,9 +164,9 @@ def main():
         optim_step += 1
         log_dict = metrics.calculate()
         log_dict.update(model.calculate_metric(is_training=True))
-        logger.info(f"[Step {optim_step}/{MAX_STEPS}] {log_dict}")
+        logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}')
 
-    logger.info(f"Training completed. optim_steps={optim_step}")
+    logger.info(f'Training completed. optim_steps={optim_step}')
     model.save('grpo-gsm8k-checkpoint')
 
 if __name__ == '__main__':

src/twinkle/infra/_ray/resource_manager.py

@@ -86,7 +86,11 @@ def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[De
         for i in range(self.nnodes):
             # TODO not accurate, because placement_group cannot distribute to node same ordered with self.nodes
             node_idx = self.min_node_idx + i if device_type != 'CPU' else i
-            node = self.nodes[node_idx]
+            try:
+                node = self.nodes[node_idx]
+            except IndexError:
+                # node_idx may not be continuous
+                node = self.nodes[0]
             node_cpu = int(node['Resources']['CPU'])
             if device_type != 'CPU':
                 bundles.append({device_type: nproc_per_node, 'CPU': max(node_cpu // 2, 1)})  # create bundles

src/twinkle/preprocessor/llm.py

@@ -1,9 +1,9 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
+import re
+
 from twinkle.data_format import Message, Trajectory
 from .base import Preprocessor
-import re
 
-from twinkle.data_format import Trajectory, Message
 
 class CompetitionMathProcessor(Preprocessor):
 
@@ -83,19 +83,18 @@ def __call__(self, row) -> Trajectory:
         ]
         return Trajectory(messages=messages, user_data=[{'target': target, 'nums': nums}])
 
+
 class GSM8KProcessor(Preprocessor):
     """Preprocessor for GSM8K dataset.
 
     GSM8K fields: question (str), answer (str ending with '#### <number>')
     Extracts the ground truth number and stores it in user_data for reward.
     """
 
-    system_prompt = (
-        "You are a helpful math assistant. Solve the problem step by step. "
-        "Show your reasoning in <think> </think> tags, then give the final "
-        "numerical answer after ####.\n"
-        "For example:\n<think> ... reasoning ... </think>\n#### 42"
-    )
+    system_prompt = ('You are a helpful math assistant. Solve the problem step by step. '
+                     'Show your reasoning in <think> </think> tags, then give the final '
+                     'numerical answer after ####.\n'
+                     'For example:\n<think> ... reasoning ... </think>\n#### 42')
 
     def extract_ground_truth(answer_str: str) -> str:
         """Extract the number after '####' from GSM8K answer."""

src/twinkle/reward/__init__.py

@@ -2,5 +2,5 @@
 from .base import Reward
 from .count_down_accuracy import CountDownAccuracy
 from .format_reward import FormatReward
-from .math_reward import MathReward
 from .gsm8k import GSM8KAccuracyReward, GSM8KFormatReward
+from .math_reward import MathReward

src/twinkle/reward/gsm8k.py

@@ -1,5 +1,6 @@
-from typing import List, Dict, Any
 import re
+from typing import Any, Dict, List
+
 from twinkle.reward.base import Reward
 
 
@@ -64,9 +65,7 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]:
                 if msg.get('role') == 'assistant':
                     completion = msg.get('content', '')
                     break
-            has_think = bool(
-                re.search(r'<think>.*?</think>', completion, re.DOTALL)
-            )
+            has_think = bool(re.search(r'<think>.*?</think>', completion, re.DOTALL))
             has_answer = bool(re.search(r'####\s*[\-\d,\.]+', completion))
             rewards.append(1.0 if (has_think and has_answer) else 0.0)
         return rewards