update

Author: Yunnglin

cookbook/client/tinker/gsm8k_grpo.py

@@ -17,7 +17,7 @@
 #
 # The server must be running first (see server.py and server_config.yaml).
 # Requires both model and sampler services to be configured.
-
+import os
 import gc
 import re
 import numpy as np
@@ -38,12 +38,12 @@
 logger = get_logger()
 
 # ========== Configuration ==========
-BASE_MODEL = 'Qwen/Qwen2.5-3B-Instruct'
+BASE_MODEL = 'Qwen/Qwen2.5-7B-Instruct'
 NUM_GENERATIONS = 4
-MAX_NEW_TOKENS = 2048
-LEARNING_RATE = 1e-5
+MAX_NEW_TOKENS = 1024
+LEARNING_RATE = 1e-4
 MAX_STEPS = 100
-BATCH_SIZE = 2
+BATCH_SIZE = 4
 TEMPERATURE = 1.0
 SYNC_INTERVAL = 1       # Save weights for sampler every N steps
 LORA_RANK = 8
@@ -60,7 +60,7 @@
 USE_SWANLAB = True
 if USE_SWANLAB:
     import swanlab
-    swanlab.login(api_key=os.environ['SWANLAB_API_KEY'], save=True)
+    swanlab.login(api_key=os.environ['SWANLAB_API_KEY'])
     swanlab.init(project="twinkle-gsm8k", config={
         'model_id': BASE_MODEL,
     })
@@ -363,8 +363,8 @@ def main():
             input_tokens = prompt_ids + sampled_tokens[:-1]
             target_tokens = [0] * ob_len + sampled_tokens
             weights = [0] * ob_len + [1] * len(sampled_tokens)
-            padded_advantages = [advantage] * len(sampled_tokens)
-            padded_logprobs = logprobs
+            padded_advantages = [0.0] * ob_len + [advantage] * len(sampled_tokens)
+            padded_logprobs = [0.0] * ob_len + logprobs
 
             datum = types.Datum(
                 model_input=types.ModelInput.from_ints(input_tokens),
@@ -393,7 +393,8 @@ def main():
 
         # ========== 7. Log ==========
         log_dict = metrics.calculate()
-        log_dict['train/loss_per_token'] = float(avg_loss)
+        if optim_result.metrics:
+            log_dict.update(optim_result.metrics)
         log_dict['train/frac_reward_zero_std'] = frac_zero_std
         log_dict['train/num_training_samples'] = len(training_data)
         logger.info(f"Step {step}: {log_dict}")

cookbook/client/tinker/megatron/server.py

@@ -15,7 +15,7 @@
 
 # Resolve the path to server_config.yaml relative to this script's location
 file_dir = os.path.abspath(os.path.dirname(__file__))
-config_path = os.path.join(file_dir, 'server_config_3b.yaml')
+config_path = os.path.join(file_dir, 'server_config.yaml')
 
 # Launch the Twinkle server — this call blocks until the server is shut down
 launch_server(config_path=config_path)

…nt/tinker/megatron/server_config_3b.yaml …nt/tinker/megatron/server_config_7b.yamlcookbook/client/tinker/megatron/server_config_3b.yaml renamed to cookbook/client/tinker/megatron/server_config_7b.yaml cookbook/client/tinker/megatron/server_config_3b.yaml renamed to cookbook/client/tinker/megatron/server_config_7b.yaml

@@ -33,23 +33,24 @@ applications:
 
   # 2. Model Service (commented out) - Would host the base model for training.
   #    Uncomment and configure if you need a training model worker.
-  - name: models-Qwen2.5-3B-Instruct
-    route_prefix: /api/v1/model/Qwen/Qwen2.5-3B-Instruct
+  - name: models-Qwen2.5-7B-Instruct
+    route_prefix: /api/v1/model/Qwen/Qwen2.5-7B-Instruct
     import_path: model
     args:
-      use_megatron: true                          # Use HuggingFace Transformers backend
-      model_id: "ms://Qwen/Qwen2.5-3B-Instruct" # ModelScope model identifier
+      use_megatron: true                          
+      model_id: "ms://Qwen/Qwen2.5-7B-Instruct" # ModelScope model identifier
+      max_length: 10240
       nproc_per_node: 2                            # Number of GPU processes per node
       device_group:
         name: model
-        ranks: [0, 1]                              # GPU rank indices
+        ranks: [0,1]                              # GPU rank indices
         device_type: cuda
       device_mesh:
         device_type: cuda
-
+        dp_size: 2
       queue_config:
         rps_limit: 100                             # Max requests per second
-        tps_limit: 10000                           # Max tokens per second
+        tps_limit: 100000                           # Max tokens per second
       adapter_config:
         per_token_adapter_limit: 30                # Max concurrent LoRA adapters
         adapter_timeout: 1800                      # Seconds before idle adapter unload
@@ -68,24 +69,28 @@ applications:
 
   # 3. Sampler Service - Runs inference / sampling using vLLM engine
   #    Used for generating text from the model (e.g., evaluating LoRA results).
-  - name: sampler-Qwen2.5-3B-Instruct
-    route_prefix: /api/v1/sampler/Qwen/Qwen2.5-3B-Instruct
+  - name: sampler-Qwen2.5-7B-Instruct
+    route_prefix: /api/v1/sampler/Qwen/Qwen2.5-7B-Instruct
     import_path: sampler
     args:
-      model_id: "ms://Qwen/Qwen2.5-3B-Instruct"   # ModelScope model identifier
-      nproc_per_node: 1               # Number of GPU processes per node
+      model_id: "ms://Qwen/Qwen2.5-7B-Instruct"   # ModelScope model identifier
+      nproc_per_node: 2               # Number of GPU processes per node
       sampler_type: vllm              # Inference engine: 'vllm' (fast) or 'torch' (TorchSampler)
       engine_args:                    # vLLM engine-specific settings
         max_model_len: 4096           # Maximum sequence length the engine supports
-        gpu_memory_utilization: 0.7   # Fraction of GPU memory to use (0.0-1.0)
+        gpu_memory_utilization: 0.5   # Fraction of GPU memory to use (0.0-1.0)
         enable_lora: true             # Allow loading LoRA adapters during inference
+        logprobs_mode: processed_logprobs # Logprobs mode for sampling results
       device_group:                   # Logical device group for the sampler
         name: sampler
-        gpus_per_worker: 1
-        ranks: [0]                    # GPU rank indices to use
+        ranks: [2]                    # GPU rank indices to use
         device_type: cuda
       device_mesh:
         device_type: cuda
+        dp_size: 1
+      queue_config:
+        rps_limit: 100                             # Max requests per second
+        tps_limit: 100000                           # Max tokens per second
     deployments:
       - name: SamplerManagement
         autoscaling_config:

cookbook/client/tinker/transformer/server.py

@@ -7,8 +7,6 @@
 
 import os
 
-# Enable Ray debug mode for verbose logging during development
-# os.environ['RAY_DEBUG'] = '1'
 os.environ['TWINKLE_TRUST_REMOTE_CODE'] = '0'
 
 from twinkle.server import launch_server

cookbook/client/tinker/transformer/server_config.yaml

@@ -33,16 +33,17 @@ applications:
 
   # 2. Model Service (commented out) - Would host the base model for training.
   #    Uncomment and configure if you need a training model worker.
-  - name: models-Qwen2.5-3B-Instruct
-    route_prefix: /api/v1/model/Qwen/Qwen2.5-3B-Instruct
+  - name: models-Qwen2.5-7B-Instruct
+    route_prefix: /api/v1/model/Qwen/Qwen2.5-7B-Instruct
     import_path: model
     args:
       use_megatron: false                          # Use HuggingFace Transformers backend
-      model_id: "ms://Qwen/Qwen2.5-3B-Instruct" # ModelScope model identifier
+      model_id: "ms://Qwen/Qwen2.5-7B-Instruct" # ModelScope model identifier
+      max_length: 10240
       nproc_per_node: 2                            # Number of GPU processes per node
       device_group:
         name: model
-        ranks: [1,2]                              # GPU rank indices
+        ranks: [0,1]                              # GPU rank indices
         device_type: cuda
       device_mesh:
         device_type: cuda
@@ -68,12 +69,12 @@ applications:
 
   # 3. Sampler Service - Runs inference / sampling using vLLM engine
   #    Used for generating text from the model (e.g., evaluating LoRA results).
-  - name: sampler-Qwen2.5-3B-Instruct
-    route_prefix: /api/v1/sampler/Qwen/Qwen2.5-3B-Instruct
+  - name: sampler-Qwen2.5-7B-Instruct
+    route_prefix: /api/v1/sampler/Qwen/Qwen2.5-7B-Instruct
     import_path: sampler
     args:
-      model_id: "ms://Qwen/Qwen2.5-3B-Instruct"   # ModelScope model identifier
-      nproc_per_node: 1               # Number of GPU processes per node
+      model_id: "ms://Qwen/Qwen2.5-7B-Instruct"   # ModelScope model identifier
+      nproc_per_node: 2               # Number of GPU processes per node
       sampler_type: vllm              # Inference engine: 'vllm' (fast) or 'torch' (TorchSampler)
       engine_args:                    # vLLM engine-specific settings
         max_model_len: 4096           # Maximum sequence length the engine supports
@@ -82,7 +83,7 @@ applications:
         logprobs_mode: processed_logprobs # Logprobs mode for sampling results
       device_group:                   # Logical device group for the sampler
         name: sampler
-        ranks: [3]                    # GPU rank indices to use
+        ranks: [2]                    # GPU rank indices to use
         device_type: cuda
       device_mesh:
         device_type: cuda

src/twinkle/infra/_ray/resource_manager.py

@@ -71,7 +71,7 @@ def __init__(self,
             self.min_node_idx = 0
             self.nnodes = math.ceil(cpu_proc_count / ncpu_proc_per_node)
 
-        breakpoint()
+        # breakpoint()
         self.nodes = []
         for node in ray.nodes():
             # get available nodes

src/twinkle/loss/grpo.py

@@ -305,7 +305,7 @@ def __call__(
         Returns:
             loss: Scalar loss value
         """
-        breakpoint()
+        # breakpoint()
         import torch
         labels = inputs.get('labels')
         assert labels is not None, "inputs must contain 'labels'"

src/twinkle/server/tinker/common/megatron_model.py

@@ -48,6 +48,18 @@ def _collect_forward_backward_results(results):
 
     return [all_outputs, avg_loss]
 
+def _clean_metrics(metrics: dict) -> dict:
+    cleaned = {}
+    for key, value in metrics.items():
+        if isinstance(value, str):
+            import re
+            match = re.match(r'^([+-]?\d*\.?\d+)', value.strip())
+            if match:
+                cleaned[key] = float(match.group(1))
+        else:
+            cleaned[key] = value
+    return cleaned
+
 
 @remote_class(execute='all')
 class TwinkleCompatMegatronModel(_MegatronBase):
@@ -179,6 +191,13 @@ def step(self, *, adam_params: types.AdamParams, **kwargs):
         # Zero gradients
         super().zero_grad(**kwargs)
 
+
+    @remote_function(collect='first', lazy_collect=False)
+    def calculate_metric(self, is_training, **kwargs):
+        metric = super().calculate_metric(is_training, **kwargs)
+        return _clean_metrics(metric)
+
+
     @remote_function(dispatch='all', sync=True)
     def load(self, checkpoint_dir: str, **kwargs):
         """

src/twinkle/server/tinker/common/transformers_model.py

@@ -35,6 +35,18 @@ def _collect_forward_backward_results(results):
 
     return [all_outputs, avg_loss]
 
+def _clean_metrics(metrics: dict) -> dict:
+    cleaned = {}
+    for key, value in metrics.items():
+        if isinstance(value, str):
+            import re
+            match = re.match(r'^([+-]?\d*\.?\d+)', value.strip())
+            if match:
+                cleaned[key] = float(match.group(1))
+        else:
+            cleaned[key] = value
+    return cleaned
+
 
 @remote_class()
 class TwinkleCompatTransformersModel(MultiLoraTransformersModel):
@@ -102,7 +114,7 @@ def forward_backward(self, *, inputs: List[types.Datum], adapter_name: str, loss
 
         # Convert Datum to InputFeature
         input_features = datum_to_input_feature(inputs, template)
-
+        # breakpoint()
         # Forward pass
         outputs = super().forward(inputs=input_features, adapter_name=adapter_name, **kwargs)
 
@@ -139,7 +151,11 @@ def step(self, *, adam_params: types.AdamParams, **kwargs):
         # Zero gradients
         super().zero_grad(**kwargs)
 
-        return super().calculate_metric(is_training=True, **kwargs)
+    @remote_function(collect='first', lazy_collect=False)
+    def calculate_metric(self, is_training, **kwargs):
+        metric = super().calculate_metric(is_training, **kwargs)
+        return _clean_metrics(metric)
+
     @remote_function()
     def load(self, checkpoint_dir: str, **kwargs):
         """

src/twinkle/server/tinker/model.py

@@ -419,8 +419,9 @@ async def _do_optim():
                     # Touch adapter to reset inactivity counter
                     self.touch_adapter(adapter_name)
 
-                    metrics = self.model.step(adam_params=body.adam_params,
+                    self.model.step(adam_params=body.adam_params,
                                     adapter_name=adapter_name)
+                    metrics = self.model.calculate_metric(is_training=True, adapter_name=adapter_name)
                     return types.OptimStepResponse(metrics=metrics)
                 except Exception:
                     logger.error(traceback.format_exc())