Naseer-010 commited on
Commit
e31649e
Β·
verified Β·
1 Parent(s): 7296b3b

Add model card with benchmark results, architecture, and reward docs

Browse files
Files changed (1) hide show
  1. README.md +237 -3
README.md CHANGED
@@ -1,3 +1,237 @@
1
- ---
2
- license: mit
3
- ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ base_model: Qwen/Qwen3-8B
4
+ tags:
5
+ - reinforcement-learning
6
+ - grpo
7
+ - infrastructure-management
8
+ - sre
9
+ - kubernetes
10
+ - lora
11
+ - unsloth
12
+ - trl
13
+ language:
14
+ - en
15
+ pipeline_tag: text-generation
16
+ ---
17
+
18
+ # Qwen3-8B β€” GRPO Fine-tuned on DIME
19
+
20
+ **Qwen3-8B fine-tuned via Group Relative Policy Optimization (GRPO) to act as an autonomous Site-Reliability Engineer in a simulated 8-node Kubernetes cluster.**
21
+
22
+ ---
23
+
24
+ ## What is this model?
25
+
26
+ This is a **merged BF16 checkpoint** of [Qwen/Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B) after 300 steps of GRPO fine-tuning on the DIME (Distributed Infrastructure Management Environment) benchmark β€” trained in **44 minutes on a single A100-SXM4-80GB**.
27
+
28
+ The model observes per-node CPU, memory, queue depths, and tail-latency telemetry, then outputs a single `kubectl` command to maintain cluster health. It was trained from scratch with a completely redesigned reward signal after the original reward function was found to produce zero-variance advantages that blocked all gradient flow.
29
+
30
+ ---
31
+
32
+ ## Benchmark Results
33
+
34
+ **+44.2% relative improvement over zero-shot Qwen3-8B** on the 14-task DIME benchmark.
35
+
36
+ ![Benchmark comparison](fig3_benchmark.png)
37
+
38
+ | Metric | Zero-shot | Fine-tuned |
39
+ |---|---|---|
40
+ | Overall avg score | 0.394 | **0.569** (best episode) |
41
+ | connection_pool_deadlock | 0.630 | **0.976** |
42
+ | memory_leak_slow_burn | 0.990 | **0.990** |
43
+ | node_failure | 0.220 | **0.920** |
44
+ | retry_storm | 0.377 | **0.587** |
45
+ | thundering_herd | 0.393 | **0.606** |
46
+ | traffic_spike | 0.024 | **0.399** |
47
+
48
+ ---
49
+
50
+ ## The Reward Engineering Challenge
51
+
52
+ The central technical contribution is replacing a **gradient-blocking reward cliff** with a differentiable seven-component signal.
53
+
54
+ The original reward returned `r = βˆ’1000` whenever the database node failed β€” which happened within the first 3 steps of most episodes. With all rewards identical, GRPO advantages collapsed to zero and no gradient flowed.
55
+
56
+ ![Reward cliff vs bounded signal](fig1_reward_cliff.png)
57
+
58
+ The fixed reward:
59
+
60
+ $$R_{env}(s,a) = \text{clip}(\, r_{topo} + r_{shed} + r_{mem} + r_{fric} + r_{lat} + r_{up} + r_{eff},\; -5,\; +5 \,)$$
61
+
62
+ with bounded components ensuring non-zero gradient everywhere.
63
+
64
+ ---
65
+
66
+ ## DIME Environment
67
+
68
+ ![DIME architecture](fig4_dime_architecture.png)
69
+
70
+ - **8-node cluster**: node-0 is a stateful PostgreSQL DB (SPOF), nodes 1–7 are stateless workers
71
+ - **Partial observability**: telemetry dropout when `cpu_i = βˆ’1`
72
+ - **6 action types**: restart, reroute, scale_up, throttle, query_logs, no_op
73
+ - **14 failure scenarios**: traffic spikes, node failures, memory leaks, retry storms, split-brain, and more
74
+ - **Error budget**: throttle burns irreplaceable budget; the agent must be economical
75
+
76
+ ---
77
+
78
+ ## Training Dynamics
79
+
80
+ ![Training curves](fig2_training_curves.png)
81
+
82
+ - `reward_format` reached 3.0 (perfect) from step 1 β€” the model learned XML scaffold immediately
83
+ - `reward_validity` stabilised at 1.9+ β€” no invalid commands after step ~10
84
+ - `reward_env` improved steadily β€” environment physics signal dominated learning
85
+ - `clipped_ratio` stayed near 0 throughout β€” healthy PPO clip utilisation
86
+ - Total wall-clock: **44 minutes** at 8.4 s/step
87
+
88
+ ---
89
+
90
+ ## Triage Oracle and the Priority Inversion Bug
91
+
92
+ The model is guided during training by a 10-rule deterministic triage tree:
93
+
94
+ ![Triage tree](fig7_triage_tree.png)
95
+
96
+ A critical discovery: the original oracle evaluated the **Black Swan rule** (`|F(s)| β‰₯ 2 β†’ throttle(0.3)`) *before* the **DB Recovery rule** (`0 ∈ F(s) β†’ restart_node(0)`). When multiple nodes including the database were failed, the oracle prescribed `throttle` instead of `restart_node(0)`. The model faithfully learned this suboptimal policy. Fixing the priority ordering accounted for **+0.044 benchmark score** improvement.
97
+
98
+ ---
99
+
100
+ ## Failure Modes Documented
101
+
102
+ Five training configurations failed before convergence β€” each diagnosable from standard TRL metrics:
103
+
104
+ ![Failure modes](fig6_failure_modes.png)
105
+
106
+ | Run | Failure | Signal |
107
+ |---|---|---|
108
+ | vLLM on A100-40GB | SM 8.0 segfault (`compilation_config` not set) | Crash at init |
109
+ | batch=4, gen=8 | CPU-bound rewards; 126 s/step, GPU idle | `samples/sec = 0.06` |
110
+ | max_comp=256 | `<think>` blocks truncated before `</think>` | `frac_reward_zero_std = 1.0` |
111
+ | reward_envΓ—2 | 10:1 env-to-triage ratio recreated zero-variance | `zero_std β†’ 1.0` at step 119 |
112
+ | oracle inverted | Learned `throttle` in DB-failure states | Low `triage/mean` |
113
+
114
+ ---
115
+
116
+ ## Reward Components
117
+
118
+ ![Reward components](fig5_reward_components.png)
119
+
120
+ ```
121
+ R(a,s) = R_fmt [-3,+3] + R_val [-2,+2] + R_env [-5,+5] + R_tri [-0.5,+1]
122
+ ```
123
+
124
+ - **R_fmt**: XML scaffold compliance (`<reasoning>` + `<action>` tags)
125
+ - **R_val**: Syntactic kubectl parse success
126
+ - **R_env**: 7-component physics reward (topology, latency, memory, uptime, budget)
127
+ - **R_tri**: Oracle triage alignment (gentle guidance, not primary teacher)
128
+
129
+ ---
130
+
131
+ ## Quick Start
132
+
133
+ ```python
134
+ from transformers import AutoModelForCausalLM, AutoTokenizer
135
+
136
+ model = AutoModelForCausalLM.from_pretrained(
137
+ "Naseer-010/Qwen3-8B-Finetuned-DIME",
138
+ torch_dtype="bfloat16",
139
+ device_map="auto",
140
+ )
141
+ tokenizer = AutoTokenizer.from_pretrained("Naseer-010/Qwen3-8B-Finetuned-DIME")
142
+
143
+ system_prompt = """You are an autonomous SRE agent managing an 8-node Kubernetes cluster.
144
+ Node-0 is the PostgreSQL database (SPOF). Nodes 1-7 are stateless workers.
145
+
146
+ TRIAGE PRIORITY (check in order):
147
+ 1. OOM: if any node mem > 0.92 β†’ kubectl delete pod node-<i>
148
+ 2. DB RECOVERY: if node-0 in failed_nodes β†’ kubectl delete pod node-0
149
+ 3. SPLIT-BRAIN: if io_wait > 0.80 β†’ kubectl throttle ingress --rate=0.5
150
+ 4. HOT-SHARD: if one worker cpu > 0.90, others low β†’ reroute traffic
151
+ 5. RETRY STORM: if p99 > 100ms and rr > 150 β†’ kubectl throttle ingress --rate=0.4
152
+ 6. ZOMBIE NODE: if worker cpu near 0 β†’ reroute away from it
153
+ 7. BLACK SWAN: if 2+ nodes failed (DB alive) β†’ kubectl throttle ingress --rate=0.3
154
+ 8. DB STRESS: if node-0 cpu > 0.80 β†’ kubectl throttle ingress --rate=0.7
155
+ 9. SAFE SCALE: if avg worker cpu > 0.75 and budget > 20 β†’ scale up
156
+ 10. HEALTHY β†’ no_op
157
+
158
+ Output format:
159
+ <reasoning>One sentence identifying which rule applies.</reasoning>
160
+ <action>{"command": "kubectl ..."}</action>"""
161
+
162
+ obs = {
163
+ "cpu_loads": [0.45, 0.82, 0.79, 0.88, 0.75, 0.81, 0.77, 0.73],
164
+ "mem_utilizations": [0.41, 0.68, 0.71, 0.65, 0.62, 0.70, 0.66, 0.64],
165
+ "queue_lengths": [12, 45, 41, 53, 38, 44, 40, 37],
166
+ "failed_nodes": [],
167
+ "latency_ms": 187.3,
168
+ "p99_latency": 312.5,
169
+ "request_rate": 1840.0,
170
+ "io_wait": 0.12,
171
+ "error_budget": 85,
172
+ "step": 4,
173
+ "task_hint": "System is under heavy traffic load."
174
+ }
175
+
176
+ import json
177
+ messages = [
178
+ {"role": "system", "content": system_prompt},
179
+ {"role": "user", "content": f"Current system state:\n{json.dumps(obs, indent=2)}\nWhat action should be taken?"}
180
+ ]
181
+
182
+ text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
183
+ inputs = tokenizer(text, return_tensors="pt").to(model.device)
184
+
185
+ outputs = model.generate(
186
+ **inputs,
187
+ max_new_tokens=1024,
188
+ temperature=0.6,
189
+ top_p=0.95,
190
+ do_sample=True,
191
+ )
192
+ print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))
193
+ ```
194
+
195
+ Expected output:
196
+ ```
197
+ <think>
198
+ P99 latency is 312ms with request rate 1840 rps and no failed nodes.
199
+ Rule 5 (retry storm): p99 > 100ms and rr > 150 β†’ throttle at 0.4.
200
+ </think>
201
+ <reasoning>Rule 5 applies: p99 latency 312ms exceeds threshold with high request rate 1840 rps β€” throttle ingress to shed load.</reasoning>
202
+ <action>{"command": "kubectl throttle ingress --rate=0.4"}</action>
203
+ ```
204
+
205
+ ---
206
+
207
+ ## Training Details
208
+
209
+ | Parameter | Value |
210
+ |---|---|
211
+ | Base model | `Qwen/Qwen3-8B` (BF16) |
212
+ | Method | GRPO (TRL 0.24.0 + Unsloth + vLLM 0.6.3) |
213
+ | LoRA rank | 32, alpha=64, all projection layers |
214
+ | Trainable params | 1.05% (349 MB adapter) |
215
+ | Training steps | 300 |
216
+ | Batch size | 1 Γ— 4 generations = 4 completions/step |
217
+ | Learning rate | 5e-6, cosine schedule |
218
+ | Max completion length | 1024 tokens |
219
+ | GPU | A100-SXM4-80GB |
220
+ | Wall-clock time | 44 minutes |
221
+
222
+ ---
223
+
224
+ ## Citation
225
+
226
+ ```bibtex
227
+ @misc{dime2026,
228
+ title = {Fine-Tuning Language Models as Autonomous SREs via GRPO: The DIME Benchmark},
229
+ author = {Nithish Sriram and Naseer},
230
+ year = {2026},
231
+ url = {https://huggingface.co/Naseer-010/Qwen3-8B-Finetuned-DIME}
232
+ }
233
+ ```
234
+
235
+ ---
236
+
237
+ *Trained at SRM AP Β· Hackathon 2026*