f

training/train.py

@@ -716,29 +716,6 @@ def run_stage2(model, train_cfg, model_cfg, spec, out_dir, logger,
         ))
 
     if resume_from:
-        # ── temporary diagnostic: patch bnb load_state_dict to dump
-        # both current optimizer.param_groups and saved state_dict["param_groups"]
-        # right before the size check, so we can see the actual mismatch.
-        try:
-            import bitsandbytes.optim.optimizer as _bnb_opt
-            _orig_load = _bnb_opt.Optimizer8bit.load_state_dict
-            def _patched_load(opt_self, state_dict):
-                print("\n=== [DIAG] bnb load_state_dict at resume ===")
-                print(f"CURRENT optimizer.param_groups: {len(opt_self.param_groups)}")
-                for i, g in enumerate(opt_self.param_groups):
-                    shapes = [tuple(p.shape) for p in g["params"][:3]]
-                    print(f"  group {i}: {len(g['params'])} params, "
-                          f"wd={g.get('weight_decay')}, lr={g.get('lr')}, "
-                          f"first3 shapes={shapes}")
-                print(f"SAVED state_dict['param_groups']: {len(state_dict['param_groups'])}")
-                for i, g in enumerate(state_dict['param_groups']):
-                    print(f"  saved group {i}: {len(g['params'])} params, "
-                          f"wd={g.get('weight_decay')}, lr={g.get('lr')}")
-                print("=== [DIAG] end ===\n")
-                return _orig_load(opt_self, state_dict)
-            _bnb_opt.Optimizer8bit.load_state_dict = _patched_load
-        except Exception as _e:
-            logger.warning(f"[DIAG] failed to patch bnb load_state_dict: {_e}")
         trainer.train(resume_from_checkpoint=resume_from)
     else:
         trainer.train()

utils/checkpoint.py

@@ -114,9 +114,15 @@ def load_checkpoint(
             )
         if lora_dir.exists():
             from peft import PeftModel
+            # is_trainable=True is REQUIRED on resume: PEFT defaults to
+            # inference mode (requires_grad=False on all LoRA params), which
+            # would shrink the trainable set to projection-only (~5 tensors)
+            # and break optimizer state loading with a param-group size
+            # mismatch against the saved 261-tensor stage-2 optimizer.
             model.llm = PeftModel.from_pretrained(
                 model.llm.base_model.model,
                 str(lora_dir),
+                is_trainable=True,
             )
             print(f"[Checkpoint] LoRA adapters loaded ← {lora_dir}")
         else: