fix: agent studio gradient descent + mlx lora stability @codex

Author: DMontgomery40

server/retrieval/mlx_qwen3.py

@@ -129,10 +129,20 @@ def __init__(self, base: Any, *, r: int, scale: float, dropout_p: float) -> None
             self.scale = float(scale)
             self.dropout = nn.Dropout(float(dropout_p)) if float(dropout_p) > 0.0 else None
 
+            # MLX `nn.Linear` exposes `weight` as (out, in). However MLX `nn.QuantizedLinear`
+            # stores a packed `weight` matrix (out, packed_in) where packed_in = in / (32 / bits).
+            # We need the *logical* dims for LoRA shapes.
             try:
-                out_dim, in_dim = int(base.weight.shape[0]), int(base.weight.shape[1])
+                qlinear = getattr(nn, "QuantizedLinear", None)
+                if qlinear is not None and isinstance(base, qlinear):
+                    bits = int(getattr(base, "bits", 4) or 4)
+                    pack = max(1, int(32 // max(1, bits)))
+                    out_dim = int(base.weight.shape[0])
+                    in_dim = int(base.weight.shape[1]) * pack
+                else:
+                    out_dim, in_dim = int(base.weight.shape[0]), int(base.weight.shape[1])
             except Exception as e:
-                raise ValueError(f"Cannot infer Linear dims for LoRA injection (missing base.weight.shape): {e}") from e
+                raise ValueError(f"Cannot infer Linear dims for LoRA injection: {e}") from e
 
             # Standard LoRA init: A ~ N(0, 0.01), B = 0
             self.lora_A = mx.random.normal((self.r, in_dim)) * 0.01

server/training/mlx_qwen3_agent_trainer.py

@@ -531,13 +531,19 @@ def _dot_trainable(dirs: dict[str, Any]) -> float:
             proj_dirs_2 = None
 
         # Optimizer + schedule.
-        scheduler = optim.schedulers.join_schedules(
-            [
-                optim.schedulers.linear_schedule(0.0, float(lr), warmup_steps),
-                optim.schedulers.cosine_decay(float(lr), max(1, total_steps - warmup_steps)),
-            ],
-            [warmup_steps],
-        )
+        # MLX linear_schedule requires steps > 0; allow warmup_ratio=0.0 safely.
+        if warmup_steps <= 0:
+            scheduler = optim.schedulers.cosine_decay(float(lr), max(1, total_steps))
+        elif warmup_steps >= total_steps:
+            scheduler = optim.schedulers.linear_schedule(0.0, float(lr), max(1, warmup_steps))
+        else:
+            scheduler = optim.schedulers.join_schedules(
+                [
+                    optim.schedulers.linear_schedule(0.0, float(lr), warmup_steps),
+                    optim.schedulers.cosine_decay(float(lr), max(1, total_steps - warmup_steps)),
+                ],
+                [warmup_steps],
+            )
         optimizer = optim.AdamW(learning_rate=scheduler, weight_decay=0.01)
 
         def _grad_norm(tree: Any) -> float:

server/training/mlx_qwen3_trainer.py

@@ -548,13 +548,19 @@ def _dirs_dot(a: dict[str, Any], b: dict[str, Any]) -> float:
             proj_dirs_1 = None
             proj_dirs_2 = None
 
-        scheduler = optim.schedulers.join_schedules(
-            [
-                optim.schedulers.linear_schedule(0.0, float(lr), warmup_steps),
-                optim.schedulers.cosine_decay(float(lr), max(1, total_steps - warmup_steps)),
-            ],
-            [warmup_steps],
-        )
+        # MLX linear_schedule requires steps > 0; allow warmup_ratio=0.0 safely.
+        if warmup_steps <= 0:
+            scheduler = optim.schedulers.cosine_decay(float(lr), max(1, total_steps))
+        elif warmup_steps >= total_steps:
+            scheduler = optim.schedulers.linear_schedule(0.0, float(lr), max(1, warmup_steps))
+        else:
+            scheduler = optim.schedulers.join_schedules(
+                [
+                    optim.schedulers.linear_schedule(0.0, float(lr), warmup_steps),
+                    optim.schedulers.cosine_decay(float(lr), max(1, total_steps - warmup_steps)),
+                ],
+                [warmup_steps],
+            )
         optimizer = optim.AdamW(learning_rate=scheduler, weight_decay=0.01)
 
         def loss_fn(batch: list[LabeledPair]) -> Any:

web/src/components/AgentTraining/TrainingStudio.tsx

@@ -13,12 +13,13 @@ import { agentTrainingService, type AgentTrainRunsScope } from '@/services/Agent
 import type { AgentTrainMetricEvent, AgentTrainRun, AgentTrainRunMeta, AgentTrainStartRequest, ChatRequest, ChatResponse } from '@/types/generated';
 
 import { NeuralVisualizer, type TelemetryPoint } from '@/components/RerankerTraining/NeuralVisualizer';
+import { GradientDescentViz } from '@/components/RerankerTraining/GradientDescentViz';
 import { StudioLogTerminal } from '@/components/RerankerTraining/StudioLogTerminal';
 import { RunDiff } from './RunDiff';
 import { RunOverview } from './RunOverview';
 
 type InspectorTab = 'run-hud' | 'live-metrics' | 'overview' | 'diff' | 'config' | 'debug-prompt';
-type BottomTab = 'timeline' | 'logs';
+type BottomTab = 'timeline' | 'logs' | 'gradient';
 type LayoutPreset = 'balanced' | 'focus_viz' | 'focus_logs' | 'focus_inspector';
 
 type StudioDockRenderers = {
@@ -1545,6 +1546,9 @@ export function TrainingStudio() {
           <button className="studio-tab-btn" data-active={bottomTab === 'logs'} onClick={() => setBottomTab('logs')}>
             Logs
           </button>
+          <button className="studio-tab-btn" data-active={bottomTab === 'gradient'} onClick={() => setBottomTab('gradient')}>
+            Gradient Descent
+          </button>
           {bottomTab === 'timeline' ? (
             <input className="studio-search" placeholder="Filter events by type/message" value={eventQuery} onChange={(e) => setEventQuery(e.target.value)} />
           ) : (
@@ -1598,13 +1602,17 @@ export function TrainingStudio() {
                 })}
               </div>
             </div>
+          ) : bottomTab === 'gradient' ? (
+            <div className="studio-gradient-viz" data-testid="studio-gradient-descent-viz">
+              <GradientDescentViz events={events} />
+            </div>
           ) : (
             renderLogsBody()
           )}
         </div>
       </section>
     );
-  }, [bottomTab, eventQuery, eventVirtualItems, filteredEvents, renderLogsBody, setBottomTab]);
+  }, [bottomTab, eventQuery, eventVirtualItems, events, filteredEvents, renderLogsBody, setBottomTab]);
 
   const dockRenderers = useMemo<StudioDockRenderers>(
     () => ({
@@ -1778,4 +1786,3 @@ export function TrainingStudio() {
     </section>
   );
 }
-

web/src/components/RerankerTraining/GradientDescentViz.tsx

@@ -1,5 +1,4 @@
 import { useEffect, useMemo, useRef, useState } from 'react';
-import type { RerankerTrainMetricEvent } from '@/types/generated';
 
 type ProjectionPoint = {
   ts: string;
@@ -10,24 +9,39 @@ type ProjectionPoint = {
   dy?: number;
 };
 
-function parseProjectionPoints(events: RerankerTrainMetricEvent[]): ProjectionPoint[] {
+type MetricEventLike = {
+  ts: string;
+  step?: number | null;
+  proj_x?: number | null;
+  proj_y?: number | null;
+  proj_dx?: number | null;
+  proj_dy?: number | null;
+  // Back-compat: some older emitters stuffed telemetry into metrics.
+  metrics?: Record<string, unknown> | null;
+};
+
+function parseProjectionPoints(events: MetricEventLike[]): ProjectionPoint[] {
   const points: ProjectionPoint[] = [];
+  let prev: ProjectionPoint | null = null;
   for (const ev of events) {
-    const m = ev.metrics as Record<string, unknown> | null | undefined;
-    const x = m?.proj_x;
-    const y = m?.proj_y;
+    const m = ev.metrics;
+    const x = typeof ev.proj_x === 'number' ? ev.proj_x : m?.proj_x;
+    const y = typeof ev.proj_y === 'number' ? ev.proj_y : m?.proj_y;
     if (typeof x !== 'number' || typeof y !== 'number') continue;
 
-    const dx = m?.proj_dx;
-    const dy = m?.proj_dy;
+    const dx = typeof ev.proj_dx === 'number' ? ev.proj_dx : m?.proj_dx;
+    const dy = typeof ev.proj_dy === 'number' ? ev.proj_dy : m?.proj_dy;
+    const computedDx = prev ? x - prev.x : undefined;
+    const computedDy = prev ? y - prev.y : undefined;
     points.push({
       ts: String(ev.ts),
       step: ev.step == null ? undefined : Number(ev.step),
       x,
       y,
-      dx: typeof dx === 'number' ? dx : undefined,
-      dy: typeof dy === 'number' ? dy : undefined,
+      dx: typeof dx === 'number' ? dx : computedDx,
+      dy: typeof dy === 'number' ? dy : computedDy,
     });
+    prev = points[points.length - 1];
   }
   return points;
 }
@@ -36,7 +50,7 @@ function clamp(v: number, lo: number, hi: number): number {
   return Math.max(lo, Math.min(hi, v));
 }
 
-export function GradientDescentViz({ events }: { events: RerankerTrainMetricEvent[] }) {
+export function GradientDescentViz({ events }: { events: MetricEventLike[] }) {
   const canvasRef = useRef<HTMLCanvasElement | null>(null);
   const points = useMemo(() => parseProjectionPoints(events), [events]);
 
@@ -245,4 +259,3 @@ export function GradientDescentViz({ events }: { events: RerankerTrainMetricEven
     </div>
   );
 }
-

web/src/components/RerankerTraining/TrainingStudio.tsx

@@ -18,6 +18,7 @@ import type {
   TrainingConfig,
 } from '@/types/generated';
 import { NeuralVisualizer, type TelemetryPoint } from './NeuralVisualizer';
+import { GradientDescentViz } from './GradientDescentViz';
 import { RunDiff } from './RunDiff';
 import { RunOverview } from './RunOverview';
 import { StudioLogTerminal } from './StudioLogTerminal';
@@ -32,7 +33,7 @@ type InspectorTab =
   | 'config'
   | 'debug-score';
 
-type BottomTab = 'timeline' | 'logs';
+type BottomTab = 'timeline' | 'logs' | 'gradient';
 type LayoutPreset = 'balanced' | 'focus_viz' | 'focus_logs' | 'focus_inspector';
 
 type StudioDockRenderers = {
@@ -1677,6 +1678,9 @@ export function TrainingStudio() {
           <button className="studio-tab-btn" data-active={bottomTab === 'logs'} onClick={() => setBottomTab('logs')}>
             Logs
           </button>
+          <button className="studio-tab-btn" data-active={bottomTab === 'gradient'} onClick={() => setBottomTab('gradient')}>
+            Gradient Descent
+          </button>
           {bottomTab === 'timeline' ? (
             <input
               className="studio-search"
@@ -1727,6 +1731,10 @@ export function TrainingStudio() {
                 })}
               </div>
             </div>
+          ) : bottomTab === 'gradient' ? (
+            <div className="studio-gradient-viz" data-testid="studio-gradient-descent-viz">
+              <GradientDescentViz events={events} />
+            </div>
           ) : (
             renderLogsBody()
           )}
@@ -1739,6 +1747,7 @@ export function TrainingStudio() {
     downloadLogs,
     eventQuery,
     eventVirtualItems,
+    events,
     filteredEvents,
     renderLogsBody,
     setBottomTab,

web/src/config/routes.ts

@@ -108,6 +108,7 @@ export const routes: RouteConfig[] = [
       { id: 'graph', title: 'Graph' },
       { id: 'reranker-config', title: 'Reranker' },
       { id: 'learning-ranker', title: 'Learning Ranker' },
+      { id: 'learning-agent', title: 'Learning Agent Studio' },
       { id: 'indexing', title: 'Indexing' }
     ],
     nav: { visible: true }

web/src/styles/learning-studio.css

@@ -586,6 +586,12 @@
   gap: 8px;
 }
 
+.studio-gradient-viz {
+  min-height: 0;
+  padding: 10px;
+  overflow: auto;
+}
+
 .studio-event-row {
   border: 1px solid var(--panel-border);
   border-radius: 8px;