Report SFT loss and targets without RL metrics - #518
Conversation
|
Documentation preview: https://modal-labs-training-gym--training-gym-previews-preview-r-78da35.modal.run/518/docs |
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
| if is_sft(args): | ||
| report_samples(rollout_id, args, samples) | ||
| # Slime's custom logger return value suppresses its default RL metrics. | ||
| return True |
There was a problem hiding this comment.
🟡 SFT custom rollout hooks never run
For SFT runs, log_rollout_data returns before invoking the configured custom rollout hook. User rollout-side behavior silently disappears when switching workloads.
Prompt for agents
Preserve the custom rollout logger contract for SFT in modal_training_gym/frameworks/slime/phase_reporting.py. The SFT branch must still report sampled conversations and suppress Slime's default RL metrics, but it also needs to invoke CUSTOM_ROLLOUT_LOG_FUNCTION_PATH_KEY with the same arguments used by the RL branch. Decide how the custom hook's return value interacts with the mandatory suppression, and add coverage for an SFT recipe with a configured custom_rollout_log_function.
Was this helpful? React with 👍 or 👎 to provide feedback.
bb67017 to
3e72335
Compare
3e72335 to
eb5a7ce
Compare
| anchors = [ | ||
| node | ||
| for node in ast.walk(train) | ||
| if isinstance(node, ast.Expr) | ||
| and isinstance(node.value, ast.Call) | ||
| and ast.unparse(node.value.func) == "logger.info" | ||
| and "log_dict" in ast.unparse(node) | ||
| ] | ||
| if len(anchors) != 1: | ||
| raise RuntimeError( | ||
| "SFT reporting requires one reduced train-metrics log anchor" | ||
| ) |
There was a problem hiding this comment.
🔍 Source patch constrains Slime overlays
SFT image builds require one exact logger.info AST anchor. Compatible overlays with different logging structure cannot use SFT reporting without patch updates.
Was this helpful? React with 👍 or 👎 to provide feedback.
eb5a7ce to
42d1144
Compare
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
| steps = vol_get_summary_items( | ||
| MetadataStore.TRAINING_STEPS_SUMMARY, key=training_run_id | ||
| ) | ||
| return sorted(steps or [], key=lambda step: step["step"]) |
There was a problem hiding this comment.
🟡 Saved losses vanish from history
When a summary update fails or overlaps another save, load_training_steps reads only the incomplete summary. Canonical records remain intact, but no healing path restores them. The loss chart permanently omits recorded steps.
Prompt for agents
TrainingStep.save writes the canonical training-steps record before performing a read-modify-write update of the per-run training-steps-summary file. If that second operation fails, or concurrent requests clobber one another, load_training_steps trusts the summary and never recovers the canonical record. The generic summary healer cannot currently handle per-run summary keys. Add a recovery strategy for keyed training-step summaries, such as rebuilding a run's summary from canonical records when needed, and ensure concurrent writes cannot permanently hide records. Preserve the one-summary-read fast path when the summary is complete.
Was this helpful? React with 👍 or 👎 to provide feedback.
| TRAIN_RESULTS_SUMMARY = "train-results-summary" | ||
| TRAINING_ROLLOUTS = "training-rollouts" | ||
| TRAINING_STEPS = "training-steps" | ||
| TRAINING_STEPS_SUMMARY = "training-steps-summary" |
There was a problem hiding this comment.
🟡 Deleted runs retain loss summaries
Deleting an old failed run leaves its loss summary behind. cleanup never removes this per-run store. Repeated cleanup cannot reclaim these files.
Prompt for agents
The cleanup command removes metadata associated with old failed or cancelled runs, but the new MetadataStore.TRAINING_STEPS_SUMMARY stores one file keyed by training_run_id and cleanup never deletes that file. Extend modal_training_gym/cli/cleanup.py to remove each target run's keyed training-step summary. Consider the canonical TRAINING_STEPS records in the same lifecycle so cleanup does not leave either representation behind, while respecting the incremental scope of this new summary store.
Was this helpful? React with 👍 or 👎 to provide feedback.
b6ec5ca to
d7a2037
Compare
| step = TrainingStep( | ||
| training_run_id=run_id, | ||
| step=int(metrics["train/step"]), | ||
| loss=metrics["train/loss"], | ||
| grad_norm=metrics.get("train/grad_norm"), | ||
| learning_rate=metrics.get("train/lr-pg_0"), | ||
| created_at=time.time(), |
There was a problem hiding this comment.
🔴 Non-finite gradient norms abort SFT
After an overflow skips an optimizer update, report_train_metrics validates the NaN gradient norm and raises. The reporting hook then aborts training instead of omitting that optional metric.
Learn more
Slime initializes grad_norm to NaN. When gradient preparation detects an overflow, train_one_step skips the optimizer update and returns that NaN sentinel while continuing the loop. TrainingStep forbids non-finite values for every float field, so constructing the report raises a Pydantic validation error. The injected call is synchronous and does not catch that error, which terminates the training process. Loss must remain strictly finite, but optional diagnostics can be absent when Slime cannot produce a finite value.
Example: An SFT step encounters an FP16 overflow with check_for_nan_in_loss_and_grad=False. Slime skips the update and returns grad_norm=NaN; metric reporting raises before the next step starts.
Recommended fix: Normalize non-finite optional grad_norm and learning_rate values to None before constructing TrainingStep. Keep non-finite loss values invalid.
Was this helpful? React with 👍 or 👎 to provide feedback.
| reporting._enqueue( | ||
| step.model_dump(), | ||
| url=reporting._derive_url("/api/training-steps"), | ||
| timeout_seconds=reporting._STEP_EVENT_TIMEOUT_SECONDS, | ||
| ) |
There was a problem hiding this comment.
🟡 Failed posts permanently drop SFT steps
report_train_metrics sends each completed step once, so queue overflow or a failed HTTP post permanently loses that step. A failed final post leaves SFT progress and latest loss stale after training completes.
Learn more
The shared reporting worker makes one HTTP attempt for ordinary queue entries. _worker only retries final timing records, and _enqueue silently drops entries when its bounded queue is full. Training steps have no journal or later reconstruction path. The dashboard derives SFT progress from the newest successfully received step, so losing the newest report leaves progress behind the actual optimizer.
Example: A 100-step run completes while the dashboard is unavailable for step 99. Steps 0–98 remain visible, but the run permanently reports 99/100 and the loss from step 98.
Recommended fix: Give training-step records durable retry semantics. Persist them locally before enqueueing, retry failed delivery with bounded backoff, and reconcile pending records during shutdown or checkpoint recovery. Preserve ordering or use the step key to make retries idempotent.
Was this helpful? React with 👍 or 👎 to provide feedback.
| def report_train_metrics(args, metrics: dict) -> None: | ||
| if not is_sft(args): | ||
| return | ||
| run_id = reporting._run_context(args)["training_run_id"] | ||
| if not run_id: | ||
| return | ||
| step = TrainingStep( | ||
| training_run_id=run_id, | ||
| step=int(metrics["train/step"]), | ||
| loss=metrics["train/loss"], | ||
| grad_norm=metrics.get("train/grad_norm"), | ||
| learning_rate=metrics.get("train/lr-pg_0"), | ||
| created_at=time.time(), | ||
| ) | ||
| reporting._enqueue( | ||
| step.model_dump(), | ||
| url=reporting._derive_url("/api/training-steps"), | ||
| timeout_seconds=reporting._STEP_EVENT_TIMEOUT_SECONDS, | ||
| ) |
There was a problem hiding this comment.
🔍 Description and implementation diverge
The description promises batching, journaling, outage recovery, and supervised samples. This commit only queues individual scalar records in memory, so verify the intended revision.
Was this helpful? React with 👍 or 👎 to provide feedback.
1 participant
SFT now reports reduced optimizer loss and supervised conversation samples without producing RL reward, advantage, or generation statistics. This addresses the earlier successful training run whose dashboard stopped at 81/100 steps.
Metrics are batched into 100-step chunks using existing metadata storage. A checkpoint-side journal is written before background publication and reconciled after training. Samples reuse the existing rollout storage/export, capped at two conversations every ten steps plus the final step. Adds a loss-history read endpoint and specializes
run get --verbose; no new metrics service or HTTP ingestion API.Stack: #517 → #518 → #519. Base:
omkaark/sft. Review only the diff against #517; #519 consumes this contract.Validation: the combined stack passes 980 Python tests, one skipped. Includes a deterministic 100-step outage/recovery test, retry ordering, non-finite loss rejection, patch idempotence against pinned Slime fixtures, API coverage, and target preservation. Self-review/live testing also added regressions for stale status writes overwriting loss/progress and no-op RL timing/status noise; journal commits now tolerate in-flight checkpoint writers. Lint and compile checks pass.
Live one-H100 Qwen3-0.6B smoke: all ten scalar records, 10/10 actual progress, loss 2.2542 → 0.7857, both configured sample batches retained, and timings for all ten steps with no reward/inference/weight-sync phases. Full run/checkpoint evidence is in #519. Broad GPU model-validation CI was cancelled before execution to honor the budget; snapshot CI requires an authorized environment reviewer.