Clear max_seq_len

Author: Ziems

.gitignore

@@ -55,3 +55,4 @@ mise.toml
 uv.lock
 **/wandb/
 .ruff_cache/
+*.pkl

arbor/core/config.py

@@ -10,7 +10,6 @@ class Config(BaseModel):
 
     # Basic settings
     storage_path: str = str(Path.home() / ".arbor" / "storage")
-    max_context_length: Optional[int] = None
 
     # Training settings
     accelerate_config: Optional[str] = None

arbor/integrations/dspy/arbor_provider.py

@@ -1,4 +1,3 @@
-import json
 import time
 from datetime import datetime
 from typing import TYPE_CHECKING, Any
@@ -64,13 +63,13 @@ class ArborReinforceJob(ReinforceJob):
         "max_completion_length": None,
         "gradient_checkpointing_kwargs": None,
         "bf16": False,
+        "fp16": True,
         "scale_rewards": True,
         "max_grad_norm": 1.0,
         "report_to": "none",
         "log_completions": True,
         "logging_steps": 10,
-        # By default, none is the model's max context length
-        "max_context_length": None,
+        "max_seq_len": None,
         "lora_config": None,
         "loss_type": "dapo",
         "soft_completion_penalty_length": None,
@@ -135,6 +134,7 @@ def initialize(self):
             self.DEFAULT_TRAIN_KWARGS["mask_truncated_completions"],
         )
         bf16 = self.train_kwargs.get("bf16", self.DEFAULT_TRAIN_KWARGS["bf16"])
+        fp16 = self.train_kwargs.get("fp16", self.DEFAULT_TRAIN_KWARGS["fp16"])
         scale_rewards = self.train_kwargs.get(
             "scale_rewards", self.DEFAULT_TRAIN_KWARGS["scale_rewards"]
         )
@@ -154,8 +154,8 @@ def initialize(self):
         logging_steps = self.train_kwargs.get(
             "logging_steps", self.DEFAULT_TRAIN_KWARGS["logging_steps"]
         )
-        max_context_length = self.train_kwargs.get(
-            "max_context_length", self.DEFAULT_TRAIN_KWARGS["max_context_length"]
+        max_seq_len = self.train_kwargs.get(
+            "max_seq_len", self.DEFAULT_TRAIN_KWARGS["max_seq_len"]
         )
         max_steps = self.train_kwargs.get("max_steps", 500)
         num_training_gpus = self.train_kwargs.get("num_training_gpus", 1)
@@ -186,14 +186,14 @@ def initialize(self):
                 "soft_completion_penalty_length": soft_completion_penalty_length,
                 "mask_truncated_completions": mask_truncated_completions,
                 "bf16": bf16,
+                "fp16": fp16,
                 "scale_rewards": scale_rewards,
                 "gradient_checkpointing_kwargs": gradient_checkpointing_kwargs,
                 "max_grad_norm": max_grad_norm,
                 "report_to": report_to,
                 "log_completions": log_completions,
                 "logging_steps": logging_steps,
-                # "max_context_length": max_context_length,
-                # "max_seq_len": max_context_length,
+                "max_seq_len": max_seq_len,
                 "max_steps": max_steps,
                 "loss_type": loss_type,
                 "num_training_gpus": num_training_gpus,
@@ -202,7 +202,7 @@ def initialize(self):
             },
             "inference_config": {
                 "model": finetune_model,
-                "max_context_length": max_context_length,
+                "max_seq_len": max_seq_len,
             },
             "gpu_config": {
                 "type": "multi",
@@ -214,12 +214,21 @@ def initialize(self):
         }
         url = urljoin(api_base, "fine_tuning/grpo/initialize")
         headers = {"Content-Type": "application/json"}
-        response = requests.post(url=url, headers=headers, json=data)
-        print(json.dumps(response.json(), indent=2))
-        response.raise_for_status()
-        response = response.json()
-        self.lm.model = ArborProvider._add_provider_prefix(response["current_model"])
-        self.provider_job_id = response.get("job_id")
+        try:
+            response = requests.post(url=url, headers=headers, json=data)
+            response.raise_for_status()
+            response = response.json()
+            self.lm.model = ArborProvider._add_provider_prefix(
+                response["current_model"]
+            )
+            self.provider_job_id = response.get("job_id")
+        except KeyboardInterrupt:
+            print(
+                "Keyboard interrupt received. Stopping reinforcement learning job initialization."
+            )
+        except Exception as err:
+            print(f"Error initializing reinforcement learning job: {err}")
+            raise err
 
     def _run_grpo_step_one_group(
         self,

arbor/server/api/schemas.py

@@ -179,7 +179,7 @@ class InferenceLaunchOwner(StrictBaseModel):
 
 class InferenceLaunchRequest(StrictBaseModel):
     model: str
-    max_context_length: Optional[int] = None
+    max_seq_len: Optional[int] = None
     num_gpus: Optional[int] = 1
     owner: Optional[InferenceLaunchOwner] = None
     log_file_path: Optional[str] = None

arbor/server/services/jobs/inference_job.py

@@ -106,8 +106,8 @@ def launch(
         vllm_module = "arbor.server.services.inference.vllm_serve"
         command = f"{sys.executable} -m {vllm_module} --model {self.launched_model_name} --port {self.port} --gpu-memory-utilization 0.9 --tensor-parallel-size {n_gpus} --enable_prefix_caching --disable-log-stats"
 
-        if launch_config.max_context_length:
-            command += f" --max_model_len {launch_config.max_context_length}"
+        if launch_config.max_seq_len:
+            command += f" --max_model_len {launch_config.max_seq_len}"
 
         logger.info(f"Running command: {command}")
 

arbor/server/services/jobs/inference_launch_config.py

@@ -4,7 +4,7 @@
 
 @dataclass
 class InferenceLaunchConfig:
-    max_context_length: Optional[int] = None
+    max_seq_len: Optional[int] = None
     gpu_ids: Optional[list[int]] = None
     num_gpus: Optional[int] = 1
     grpo_job_id: Optional[str] = None

arbor/server/services/managers/inference_manager.py

@@ -84,7 +84,7 @@ def launch_from_request(
                 return existing_job
 
         launch_config = InferenceLaunchConfig(
-            max_context_length=launch_request.max_context_length,
+            max_seq_len=launch_request.max_seq_len,
             gpu_ids=list(preallocated_gpu_ids) if preallocated_gpu_ids else None,
             num_gpus=launch_request.num_gpus,
             log_file_path=launch_request.log_file_path,

examples/checkpointing.py

@@ -65,7 +65,7 @@ def initialize_grpo(
         "trainer_config": trainer_config,
         "inference_config": {
             "model": model,
-            "max_context_length": 2048,
+            "max_seq_len": 2048,
         },
         "gpu_config": {
             "type": "multi",

examples/inference_swap.py

@@ -18,13 +18,13 @@
 
 
 def launch_model(
-    model_name: str, *, num_gpus: int = 1, max_context_length: int | None = None
+    model_name: str, *, num_gpus: int = 1, max_seq_len: int | None = None
 ) -> str:
     """Launch the inference server for a model via Arbor's /launch endpoint."""
 
     payload: dict[str, object] = {"model": model_name, "num_gpus": num_gpus}
-    if max_context_length is not None:
-        payload["max_context_length"] = max_context_length
+    if max_seq_len is not None:
+        payload["max_seq_len"] = max_seq_len
 
     response = requests.post(f"{chat_base}/launch", json=payload)
     response.raise_for_status()

examples/unique_chars_grpo.py

@@ -69,7 +69,7 @@ def _unique_letter_reward(input, pred, trace=None) -> float:
     "max_grad_norm": 1.0,
     "report_to": "wandb",
     "log_completions": True,
-    "max_context_length": None,
+    "max_seq_len": None, #defaults to 2048
     "max_steps": 1000,
     "logging_steps": 10,
     "loss_type": "dapo", #"dr_grpo",