convaiinnovations/flux-test-time-training
0
1import random
2import logging
3import os
4import gc
5
6# Optimize CUDA memory allocation to reduce fragmentation
7os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
8
9import torch
10import torch.nn as nn
11import torch.optim as optim
12from modeling_physics_rl import PhysicsModel, Config
13
14class StratifiedReplayBuffer:
15 """
16 Stores memories by Concept ID (or just generic 'user_taught') to ensure we sample DIVERSE history.
17 """
18 def __init__(self):
19 self.memory = {} # { "concept_id": [ {prompt, answer}, ... ] }
20 # Pre-fill with Anchor Memories to prevent Cold-Start Catastrophic Forgetting
21 self._add_anchor_memories()
22
23 def _add_anchor_memories(self):
24 anchors = [
25 ("What is gravity?", "Gravity is a fundamental interaction which causes mutual attraction between all things with mass or energy."),
26 ("Hello", "Hello! How can I help you today?"),
27 ("What is AI?", "Artificial Intelligence (AI) refers to the simulation of human intelligence in machines."),
28 ("Define thermodynamics.", "Thermodynamics is a branch of physics that deals with heat, work, and temperature, and their relation to energy, entropy, and the physical properties of matter."),
29 ("Who are you?", "I am a large language model, trained by Google.")
30 ]
31 self.memory["anchor"] = [{"prompt": q, "answer": a} for q, a in anchors]
32 print(f" ⚓ Added {len(anchors)} General Knowledge Anchors to Replay Buffer.")
33
34 def add(self, concept_id, prompt, answer):
35 if concept_id not in self.memory:
36 self.memory[concept_id] = []
37 self.memory[concept_id].append({"prompt": prompt, "answer": answer})
38
39 def sample_stratified(self, current_concept_id, n_per_concept=1):
40 batch = []
41 past_concepts = [cid for cid in self.memory.keys() if cid != current_concept_id]
42 if not past_concepts: return []
43 for cid in past_concepts:
44 samples = random.sample(self.memory[cid], min(len(self.memory[cid]), n_per_concept))
45 batch.extend(samples)
46 return batch
47
48
49class ContinuousLearningSession:
50 def __init__(self):
51 print("🧠 Initializing Continuous Learning Session...")
52
53 # 1. Load Model
54 self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
55 print(f" 🚀 Using Device: {self.device}")
56
57 self.model = PhysicsModel()
58 self.model.to(self.device) # Force move to GPU
59
60 # 2. Load Pre-trained Weights
61 self._load_weights()
62
63 # 3. Setup Online Optimizer
64 # Update BOTH Controller AND Flux Adapters for true adaptation
65 trainable_params = [
66 {'params': self.model.controller.parameters(), 'lr': 1e-3}, # Fast adaptation
67 ]
68
69 # Also update the Flux Adapters' modulation projection
70 for layer in self.model.flux_layers:
71 trainable_params.append({'params': layer.modulation_proj.parameters(), 'lr': 5e-4})
72
73 self.optimizer = optim.AdamW(trainable_params)
74
75 # 4. Session Memory (Context Window)
76 # This stores the "learned context" so the model remembers the session
77 self.session_context = [] # List of (input, modulation) pairs
78 self.context_modulation = None # Accumulated modulation bias
79
80 # 5. Ensure backbone is frozen, but Controller & Adapters are TRAINABLE
81 for p in self.model.llm.parameters():
82 p.requires_grad = False
83
84 print(" 🔧 Unfreezing Controller & Flux Adapters...")
85 for p in self.model.controller.parameters():
86 p.requires_grad = True
87
88 if isinstance(self.model.flux_layers, list):
89 for layer in self.model.flux_layers:
90 for p in layer.parameters():
91 p.requires_grad = True
92 else:
93 for p in self.model.flux_layers.parameters():
94 p.requires_grad = True
95
96 # 3. Setup Online Optimizer
97 # Update BOTH Controller AND Flux Adapters for true adaptation
98 controller_params = list(self.model.controller.parameters())
99 if isinstance(self.model.flux_layers, torch.nn.ModuleList) or isinstance(self.model.flux_layers, torch.nn.Sequential):
100 adapter_params = list(self.model.flux_layers.parameters())
101 else:
102 # If it's a python list
103 adapter_params = [p for layer in self.model.flux_layers for p in layer.parameters()]
104
105 # Switch back to Adam (Better convergence, relying on GC/Env for memory safety)
106 self.optimizer = optim.Adam(controller_params + adapter_params, lr=1e-4)
107
108 self.model.train() # Enable gradients for Controller/Adapters
109
110 # 6. Initialize Replay Buffer & Drift Anchor
111 self.replay_buffer = StratifiedReplayBuffer()
112 self.initial_controller_state = {k: v.clone() for k, v in self.model.controller.state_dict().items()}
113
114 print(" ✅ Ready for Interactive Continuous Learning (Powered by Replay Buffer)!")
115
116 def _load_weights(self):
117 """Load pre-trained weights from various possible locations."""
118 search_paths = [
119 ".",
120 "/kaggle/input/worldmodels/physics_model",
121 "/kaggle/working/physics_model"
122 ]
123
124 for path in search_paths:
125 controller_path = os.path.join(path, "final_physics_controller.pt")
126 if os.path.exists(controller_path):
127 print(f" Loading weights from {path}...")
128 self.model.controller.load_state_dict(
129 torch.load(controller_path, map_location=self.device)
130 )
131
132 # Load WALT
133 walt_path = os.path.join(path, "final_walt_head.pt")
134 if os.path.exists(walt_path):
135 self.model.walt.load_state_dict(
136 torch.load(walt_path, map_location=self.device)
137 )
138
139 # Load Adapters
140 adapter_path = os.path.join(path, "final_liquid_adapters.pt")
141 if os.path.exists(adapter_path):
142 adapter_states = torch.load(adapter_path, map_location=self.device)
143 for layer, state in zip(self.model.flux_layers, adapter_states):
144 layer.load_state_dict(state)
145 print(" ✅ Loaded Flux Adapters.")
146
147 return
148
149 print(" ⚠️ No pre-trained weights found. Using random initialization.")
150
151
152 # def _get_context_modulation(self):
153 # """
154 # Compute a modulation bias from session history.
155 # This allows the model to "remember" previous physics context.
156 # """
157 # if not self.session_context:
158 # return None
159
160 # # Average the modulations from recent context (last 3 interactions)
161 # recent = self.session_context[-3:]
162 # mods = [m for _, m in recent if m is not None]
163
164 # if not mods:
165 # return None
166
167 # # Stack and average
168 # stacked = torch.stack(mods)
169 # return stacked.mean(dim=0)
170
171 def predict(self, user_input: str):
172 """
173 Generate a response using the current Controller & Flux Adapters.
174 Pure Inference: No context history, just the current weights.
175 """
176 self.model.eval()
177
178 full_prompt = f"User: {user_input}\nModel:"
179 inputs = self.model.tokenizer(full_prompt, return_tensors="pt").to(self.device)
180
181 # 1. Generate Modulation (Based strictly on CURRENT input)
182 with torch.no_grad():
183 h_init = self.model.get_embeddings(inputs.input_ids).to(Config.DTYPE)
184
185 modulation = self.model.controller(h_init)
186
187 # 2. No Context Bias (Disabled per request)
188 # We rely solely on the weight updates from 'learn()'
189 # context_mod = self._get_context_modulation()
190 # if context_mod is not None:
191 # # Blend: 70% new, 30% context
192 # modulation = 0.7 * modulation + 0.3 * context_mod.to(modulation.device)
193
194 # 3. Apply modulation and generate
195 self.model.set_active_modulation(modulation)
196
197 out_ids = self.model.llm.generate(
198 **inputs,
199 max_new_tokens=100, # Increased for chat
200 # max_length=Config.MAX_LENGTH, # Removed as per diff
201 do_sample=True,
202 temperature=0.7, # Changed from 0.6 to 0.7
203 repetition_penalty=1.0, # Reset to default (was 1.2) to fix silence
204 pad_token_id=self.model.tokenizer.eos_token_id
205 )
206
207 response = self.model.tokenizer.decode(out_ids[0], skip_special_tokens=True)
208 response_clean = response.split("Model:")[-1].strip()
209
210 self.model.clear_modulation()
211
212 return response_clean, modulation.detach()
213
214 def _generate_synthetic_data(self, question, answer, num_variations=3):
215 """
216 Uses the frozen Base LLM to generate diverse variations of the training example.
217 This turns One-Shot Learning into Few-Shot Learning (Synthetic Data Augmentation).
218 """
219 print(" ✨ Generating synthetic training data (Self-Distillation)...")
220
221 # 1. Disable adapters/modulation to get clean English capability
222 self.model.clear_modulation()
223 self.model.eval()
224
225 prompt = (
226 f"Original Question: {question}\n"
227 f"Original Answer: {answer}\n\n"
228 f"Task: Rewrite the above Question and Answer pair in {num_variations} different styles (e.g. simple, formal, detailed). "
229 f"Keep the facts exactly the same.\n"
230 f"Output format:\n"
231 f"Q1: ...\n"
232 f"A1: ...\n"
233 f"Q2: ...\n"
234 f"A2: ...\n"
235 f"Start now:"
236 )
237
238 inputs = self.model.tokenizer(prompt, return_tensors="pt").to(self.device)
239
240 with torch.no_grad():
241 out_ids = self.model.llm.generate(
242 **inputs,
243 max_new_tokens=256,
244 do_sample=True,
245 temperature=0.7
246 )
247
248 raw_text = self.model.tokenizer.decode(out_ids[0], skip_special_tokens=True)
249 # Parse the output (Simple heuristic parsing)
250 variations = [{"q": question, "a": answer}] # Always include original
251
252 current_q = None
253 for line in raw_text.split('\n'):
254 line = line.strip()
255 if line.startswith("Q") and ":" in line:
256 current_q = line.split(":", 1)[1].strip()
257 elif line.startswith("A") and ":" in line and current_q:
258 current_a = line.split(":", 1)[1].strip()
259 # Validation: Ensure neither Q nor A is empty or garbage
260 if current_q and current_a and "..." not in current_q and "..." not in current_a:
261 variations.append({"q": current_q, "a": current_a})
262 current_q = None
263
264 # Cleanup Memory
265 del inputs, out_ids
266 torch.cuda.empty_cache()
267
268 # Fallback: If synthetic generation failed, duplicate original
269 if len(variations) == 1:
270 print(" ⚠️ Synthetic generation failed to produce valid format. Duplicating original.")
271 variations.append({"q": question, "a": answer})
272
273 print(f" ✨ Generated {len(variations)-1} synthetic variations.")
274 for i, v in enumerate(variations):
275 print(f" [{i}] Q: {v['q'][:30]}... A: {v['a'][:30]}...")
276
277 return variations
278
279 def learn(self, user_input: str, correct_answer: str, concept_id: str = "general"):
280 """
281 Robust Learning: Updates weights using the new example + Replay Buffer.
282 Runs specific number of steps (plasticity) while anchoring to past (stability).
283 """
284 print("\n 🧠 Starting Robust Adaptation Loop...")
285
286 # 0. Augment Data (Synthetic Variations)
287 training_batch = self._generate_synthetic_data(user_input, correct_answer)
288
289 # 1. Add new knowledge to Buffer
290 self.replay_buffer.add(concept_id, user_input, correct_answer)
291
292 # Force cleanup before training to prevent OOM
293 gc.collect()
294 torch.cuda.empty_cache()
295
296 # 2. Training Loop (Micro-Epochs)
297 # 2. Training Loop (Micro-Epochs)
298 steps = 20 # Reduced to 20 (Safe limit for strong replay)
299
300 for step in range(steps):
301 self.optimizer.zero_grad()
302 total_loss = 0
303
304 # --- A. Current Task (Random Sample from Synthetic Batch) ---
305 # Pick a random variation to train on this step
306 example = random.choice(training_batch)
307
308 # Append EOS so model knows when to STOP talking
309 full_text = f"User: {example['q']}\nModel: {example['a']}{self.model.tokenizer.eos_token}"
310 inputs_train = self.model.tokenizer(full_text, return_tensors="pt", max_length=Config.MAX_LENGTH, truncation=True).to(self.device)
311
312 h_train = self.model.get_embeddings(inputs_train.input_ids).to(Config.DTYPE)
313 mod_pred = self.model.controller(h_train)
314 logits = self.model(inputs_train.input_ids, forced_modulation=mod_pred)
315
316 shift_logits = logits[..., :-1, :].contiguous()
317 shift_labels = inputs_train.input_ids[..., 1:].contiguous()
318 task_loss = torch.nn.functional.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
319
320 total_loss += task_loss * 1.0
321
322 # --- B. Replay (Stability) ---
323 past_memories = self.replay_buffer.sample_stratified(concept_id, n_per_concept=2)
324 if past_memories:
325 for mem in past_memories:
326 full_replay = f"User: {mem['prompt']}\nModel: {mem['answer']}"
327 inputs_replay = self.model.tokenizer(full_replay, return_tensors="pt", max_length=Config.MAX_LENGTH, truncation=True).to(self.device)
328
329 h_rep = self.model.get_embeddings(inputs_replay.input_ids).to(Config.DTYPE)
330 mod_rep = self.model.controller(h_rep)
331 logits_rep = self.model(inputs_replay.input_ids, forced_modulation=mod_rep)
332
333 s_log = logits_rep[..., :-1, :].contiguous()
334 s_lab = inputs_replay.input_ids[..., 1:].contiguous()
335 loss_rep = torch.nn.functional.cross_entropy(s_log.view(-1, s_log.size(-1)), s_lab.view(-1))
336
337 # Weight Replay EQUAL (1.0) to task to enforce stability
338 total_loss += loss_rep * 1.0
339
340 # --- C. Anti-Drift (Crucial for TTT) ---
341 # Penalize deviation from original weights to prevent "Model Collapse"
342 drift_loss = 0
343 for name, param in self.model.controller.named_parameters():
344 drift_loss += torch.sum((param - self.initial_controller_state[name].to(self.device)) ** 2)
345 total_loss += drift_loss * 10.0 # Very Strong anchor (was 1.0)
346
347 total_loss.backward()
348
349 # Debug Gradients
350 total_norm = 0.0
351 for p in self.model.controller.parameters():
352 if p.grad is not None:
353 total_norm += p.grad.data.norm(2).item() ** 2
354 total_norm = total_norm ** 0.5
355
356 self.optimizer.step()
357
358 if (step+1) % 10 == 0:
359 print(f" Step {step+1}: Loss {total_loss.item():.4f} | Grad Norm: {total_norm:.4f}")
360
361 # Early Stopping (Prevent Overfitting)
362 if total_loss.item() < 0.005:
363 print(f" ✅ Converged early at step {step+1} (Loss < 0.005)")
364 break
365
366 # 3. Store context (DISABLED)
367 # self.session_context.append((user_input, mod_pred.detach()))
368 self.model.clear_modulation()
369
370 print(" ✅ Adaptation Complete. Weights Updated.")
371 return total_loss.item()
372
373 def save_weights(self, suffix="session"):
374 """Save the updated weights after a learning session."""
375 print(" 💾 Saving updated weights...")
376 torch.save(self.model.controller.state_dict(), f"controller_{suffix}.pt")
377
378 adapter_states = [l.state_dict() for l in self.model.flux_layers]
379 torch.save(adapter_states, f"adapters_{suffix}.pt")
380
381 print(f" ✅ Saved to controller_{suffix}.pt and adapters_{suffix}.pt")
382
383 def run(self):
384 """Main interactive loop."""
385 print("\n" + "="*60)
386 print(" 🧪 CONTINUOUS LEARNING LAB")
387 print(" Commands:")
388 print(" - Ask any physics question")
389 print(" - Type 'wrong' if the answer is incorrect")
390 print(" - Type 'save' to save updated weights")
391 print(" - Type 'exit' to quit")
392 print("="*60)
393
394 while True:
395 try:
396 user_input = input("\nUSER: ").strip()
397 except (EOFError, KeyboardInterrupt):
398 break
399
400 if not user_input:
401 continue
402 if user_input.lower() in ['exit', 'quit']:
403 break
404 if user_input.lower() == 'save':
405 self.save_weights()
406 continue
407
408 # Generate prediction
409 response, modulation = self.predict(user_input)
410 mod_norm = modulation.norm().item()
411
412 print(f"MODEL: {response}")
413 print(f" [Modulation Norm: {mod_norm:.2f}]")
414
415 # Feedback loop
416 try:
417 feedback = input(" (Enter=correct, 'wrong'=teach): ").strip().lower()
418 except (EOFError, KeyboardInterrupt):
419 break
420
421 if feedback == "wrong":
422 try:
423 truth = input(" CORRECT ANSWER: ").strip()
424 # topic = input(" TOPIC ID (e.g. 'gravity', 'thermo'): ").strip()
425 topic = "general" # Defaulting as requested
426 except (EOFError, KeyboardInterrupt):
427 break
428
429 if truth:
430 # Pass the topic to learn so it can index it correctly
431 self.learn(user_input, truth, topic)
432 # Store correct modulation in context (DISABLED)
433 # self.session_context.append((user_input, modulation))
434 else:
435 # Correct answer - store in context for future reference (DISABLED)
436 # self.session_context.append((user_input, modulation))
437 print(" 👍 Perfect! (No update needed)")
438
439 print("\n👋 Session ended.")
440
441 # Offer to save
442 try:
443 save = input(" Save updated weights? (y/n): ").strip().lower()
444 if save == 'y':
445 self.save_weights()
446 except:
447 pass
448
449
450if __name__ == "__main__":
451 session = ContinuousLearningSession()
452 session.run()
453