Option B: --mix-scale, trackio-before-config fix

train_securecoder.py

@@ -687,6 +687,7 @@ def init_trackio(args):
         log.info("trackio initialised (project=%s space=%s)", args.trackio_project, args.trackio_space)
     except Exception as exc:  # noqa: BLE001 - monitoring must never kill training
         log.warning("trackio init failed (%s); continuing without it", exc)
+        args.report_to = "none"
 
 
 def train(args, model, tokenizer, records: list[dict]):
@@ -700,8 +701,8 @@ def train(args, model, tokenizer, records: list[dict]):
     log.info("train rows=%d eval rows=%d", len(train_ds), len(eval_ds) if eval_ds else 0)
 
     steps_per_epoch = len(train_ds) // max(args.batch_size * args.grad_accum, 1)
-    cfg = build_sft_config(args, eval_ds is not None, steps_per_epoch)
     init_trackio(args)
+    cfg = build_sft_config(args, eval_ds is not None, steps_per_epoch)
 
     trainer = SFTTrainer(
         model=model,
@@ -787,6 +788,8 @@ def parse_args(argv=None):
     p.add_argument("--learning-rate", type=float, default=2e-4)
     p.add_argument("--lr-scheduler", default="cosine")
     p.add_argument("--num-epochs", type=float, default=1.0)
+    p.add_argument("--mix-scale", type=float, default=1.0,
+                   help="multiply every source limit by this (e.g. 0.5 for a half mix)")
     p.add_argument("--max-steps", type=int, default=0, help="overrides --num-epochs when > 0")
     p.add_argument("--eval-samples", type=int, default=200, help="0 disables evaluation")
     p.add_argument("--logging-steps", type=int, default=10)
@@ -825,9 +828,15 @@ def apply_smoke(args) -> None:
 
 
 def main(argv=None) -> int:
+    global MIX
     args = parse_args(argv)
     if args.smoke:
         apply_smoke(args)
+    elif args.mix_scale != 1.0:
+        MIX = [Source(s.repo, max(50, int(s.limit * args.mix_scale)),
+                      s.kind, s.config, s.split, s.note) for s in MIX]
+        log.info("mix scaled by %.2f -> %d rows planned", args.mix_scale,
+                 sum(s.limit for s in MIX))
     token = os.environ.get("HF_TOKEN")
 
     if args.validate_only: