import json import os ALLOWED = { "learning_rate_multiplier": {0.5, 0.75, 1.0, 1.25, 1.5}, "beta_multiplier": {0.5, 1.0, 2.0}, "temperature_delta": {-0.1, 0.0, 0.1}, "max_completion_length_delta": {0}, "num_generations_delta": {0}, "early_stop": {False}, "rollback_to_best_checkpoint": {False, True}, } def no_change_decision(reason="Fixed schedule"): return { "learning_rate_multiplier": 1.0, "beta_multiplier": 1.0, "temperature_delta": 0.0, "max_completion_length_delta": 0, "num_generations_delta": 0, "early_stop": False, "rollback_to_best_checkpoint": False, "reason": reason, } def rule_decision(config, metrics, history): decision = no_change_decision("No trigger fired") reasons = [] recent = [item.get("eval_accuracy", 0.0) for item in history[-3:]] if len(recent) >= 3 and max(recent[1:]) <= recent[0] + 1e-6: decision["temperature_delta"] = 0.1 reasons.append("eval accuracy stalled") if (metrics.get("kl_mean") or 0.0) > 0.15: decision["learning_rate_multiplier"] = 0.5 decision["beta_multiplier"] = 2.0 reasons.append("KL is high") if (metrics.get("completion_length_clip_ratio") or 0.0) > 0.2: decision["max_completion_length_delta"] = 64 reasons.append("completions are frequently clipped") if (metrics.get("train_reward_std") or 0.0) < 0.02 and (metrics.get("train_reward_mean") or 0.0) < 0.8: decision["temperature_delta"] = max(decision["temperature_delta"], 0.1) reasons.append("reward diversity is near zero") decision["reason"] = "; ".join(reasons) if reasons else decision["reason"] return decision def validate_decision(decision): for key, choices in ALLOWED.items(): if key not in decision or decision[key] not in choices: raise ValueError(f"invalid action for {key}: {decision.get(key)}") decision["reason"] = str(decision.get("reason", ""))[:500] return decision def parse_llm_decision(text): marker = "FINAL_JSON:" marker_index = text.rfind(marker) if marker_index < 0: raise ValueError("LLM response is missing FINAL_JSON marker") analysis = text[:marker_index].strip() payload = text[marker_index + len(marker):].strip() if payload.startswith("```json"): payload = payload[len("```json"):].strip() elif payload.startswith("```"): payload = payload[len("```"):].strip() if payload.endswith("```"): payload = payload[:-3].strip() decision = json.loads(payload) decision["controller_analysis"] = analysis[:4000] return validate_decision(decision) CONTROLLER_METRIC_KEYS = [ "stage", "global_train_steps", "train_reward_mean", "train_reward_std", "eval_accuracy", "eval_greedy_accuracy", "eval_sampled_pass_at_1", "eval_sampled_pass_at_4", "kl_mean", "entropy_mean", "grad_norm", "last_loss", "end_learning_rate", "avg_completion_length", "completion_length_clip_ratio", "wall_clock_seconds", ] def compact_metrics(metrics): return {key: metrics.get(key) for key in CONTROLLER_METRIC_KEYS if key in metrics} def llm_decision(config, metrics, history): fallback = rule_decision(config, metrics, history) if config.get("llm_controller_mode", "mock") != "api": fallback["reason"] = "Mock LLM controller: " + fallback["reason"] return fallback try: from openai import OpenAI client = OpenAI( api_key=os.environ["OPENROUTER_API_KEY"], base_url="https://openrouter.ai/api/v1", default_headers={ "HTTP-Referer": "https://github.com/llm-zero-lite", "X-OpenRouter-Title": "LLMZero-Lite Experiment", }, ) payload = { "current_schedule": { key: config[key] for key in [ "learning_rate", "beta", "temperature", "max_completion_length", "num_generations", "per_device_train_batch_size", "gradient_accumulation_steps", ] }, "latest_metrics": compact_metrics(metrics), "recent_history": [compact_metrics(item) for item in history[-3:]], "allowed_actions": {key: sorted(values) for key, values in ALLOWED.items()}, } messages = [ { "role": "system", "content": ( "You are controlling a staged GRPO experiment. Analyze the learning dynamics, " "including reward trend and variance, evaluation accuracy, KL, entropy, gradient norm, " "completion length, and previous actions. Explain the evidence and tradeoffs in roughly " "150-400 words. Then end with FINAL_JSON: followed by exactly one JSON object and no " "text after it. Every action value must come from allowed_actions. Keep early_stop false " "so all methods receive equal compute. Include a concise reason field in the JSON." ), }, {"role": "user", "content": json.dumps(payload, indent=2)}, ] errors = [] attempt_messages = messages for attempt in range(config.get("llm_controller_max_retries", 3)): try: response = client.chat.completions.create( model=config["llm_controller_model"], temperature=0.2 if attempt == 0 else 0, max_tokens=config.get("llm_controller_max_tokens", 1200), messages=attempt_messages, timeout=90, ) content = response.choices[0].message.content or "" decision = parse_llm_decision(content) decision["controller_attempts"] = attempt + 1 return decision except Exception as exc: errors.append(f"attempt {attempt + 1}: {type(exc).__name__}: {exc}") attempt_messages = messages + [ { "role": "user", "content": ( "The previous response failed validation. Return only FINAL_JSON: followed by one " "valid JSON object using allowed_actions. Put no text after the JSON." ), }, ] raise RuntimeError("; ".join(errors)) except Exception as exc: if not config.get("llm_controller_fail_open", True): raise RuntimeError(f"LLM controller API failed: {exc}") from exc decision = no_change_decision(f"GLM failed after retries; no-change fallback: {exc}") decision["controller_failed"] = True decision["controller_analysis"] = "Controller unavailable or response invalid; preserved the current schedule." return decision def choose_decision(method, config, metrics, history): if method == "fixed_grpo": return no_change_decision() if method == "rule_controller": return rule_decision(config, metrics, history) if method == "llm_controller": return llm_decision(config, metrics, history) raise ValueError(f"unknown method: {method}") def apply_decision(config, decision): updated = dict(config) updated["learning_rate"] = min(5e-6, max(1e-7, config["learning_rate"] * decision["learning_rate_multiplier"])) updated["beta"] = min(0.2, max(0.0, config["beta"] * decision["beta_multiplier"])) updated["temperature"] = min(1.1, max(0.5, config["temperature"] + decision["temperature_delta"])) updated["max_completion_length"] = min(512, max(64, config["max_completion_length"] + decision["max_completion_length_delta"])) candidate_generations = min(16, max(2, config["num_generations"] + decision["num_generations_delta"])) generation_batch = config["per_device_train_batch_size"] * config["gradient_accumulation_steps"] updated["num_generations"] = candidate_generations if generation_batch % candidate_generations == 0 else config["num_generations"] return updated