Triage Unpack Forward Output

Triage unpack_forward_output training failures

This skill handles one failure pattern only:

tt-forge-models doesn't implement unpack_forward_output for this model.

It works on exactly one test (passed as <test-name>). Step 3's CPU triage is the real gate: if the model's forward output is already handled by an existing registry entry or loader override, the YAML is just stale and the skill skips to the YAML-update step. The YAML reason: is informational, not authoritative.

Argument: <test-name> accepted forms

Auto-detect both:

• YAML key - e.g. yolov9/pytorch-T-single_device-training. Looks like <model_dir>/<framework>-<variant>-<parallelism>-<run_mode>.
• model_info.name - e.g. pytorch_YOLOv9_T_cv_object_det_github. Composed in third_party/tt_forge_models/config.py as f"{framework}_{model}_{variant}_{task}_{source}".

Background (load into context, do not paraphrase to the user)

• Default unpack_forward_output lives in third_party/tt_forge_models/base.py and delegates to third_party/tt_forge_models/training_utils.py. The util has a registry keyed by output class name (BaseModelOutputWithPast, CausalLMOutputWithPast, CLIPOutput, ImageClassifierOutput, …).
• When the model returns a class not in the registry (a bare list/tuple/dict, or a custom dataclass), the call raises ValueError("No handler for class … exists in unpack_forward_output. Register a handler or implement custom unpack_forward_output for the specific model."). The runner then surfaces the failure with bringup_status: FAILED_FE_COMPILATION and the canonical reason: shown above.
• Two ways to fix. Pick one:
  1. Registry entry in third_party/tt_forge_models/training_utils.py - _register_attr("<OutputClassName>", "<attr>"). Preferred for HuggingFace ModelOutput dataclasses (any class that lives in transformers.models.*.modeling_*), and for any case where the loss-relevant tensor is a single attribute lookup. This is how CausalLMOutputWithPast -> logits, DPRReaderOutput -> end_logits, DepthEstimatorOutput -> predicted_depth, etc. are already wired.
  2. Per-loader override - add def unpack_forward_output(self, forward_output) to that model's ModelLoader. Required when the output has no stable class name to key on (bare list/tuple/dict), when extracting the loss-relevant tensors needs a structural transform (concat, stack, slice), or when two models share an output class but legitimately need different loss targets.
• Existing per-loader override examples to mimic when the override is the right tool:
  • third_party/tt_forge_models/yolov9/pytorch/loader.py::ModelLoader::unpack_forward_output - list-of-tensors via extract_tensors_recursive
  • third_party/tt_forge_models/centernet/pytorch/loader.py::ModelLoader::unpack_forward_output - list-of-dicts
  • third_party/tt_forge_models/clip/pytorch/loader.py::ModelLoader::unpack_forward_output - tuple
• Implication of choosing the registry path: the registry entry is global by class name, so it fixes every variant of the current model and every other model in the codebase that returns the same output class. You must therefore re-run all affected training tests, not just the one you were invoked with, and update each of their YAML entries.
• Test-runner glue: tests/runner/test_models.py parametrizes test_all_models_torch with (test_entry, parallelism, run_mode). Pytest test IDs combine in stack order: <test_entry_id>-<parallelism_id>-<run_mode_id>. The <test_entry_id> is <model_path>-<variant> (see tests/runner/utils/dynamic_loader.py:318).
• Test-runner gates on YAML status: NOT_SUPPORTED_SKIP entries are skipped (tests/runner/test_models.py:118, tests/runner/test_utils.py:746-748) and KNOWN_FAILURE_XFAIL entries are imperatively xfailed. Verification uses pytest --force-run --runxfail to bypass both without editing the YAML: --force-run (tests/runner/conftest.py:113) makes NOT_SUPPORTED_SKIP entries run their body, --runxfail makes KNOWN_FAILURE_XFAIL entries report their real outcome instead of XFAIL. Apply the YAML edit once, after observing the real outcome.

Workflow

Step 1 - Resolve <test-name> -> YAML entry + loader path

1. If <test-name> matches ^[a-z0-9_]+/.+-(single_device|data_parallel|tensor_parallel)-(training|inference)$, treat as a YAML key. Otherwise treat as a model_info.name string.
2. YAML key path: locate the entry in tests/runner/test_config/torch/test_config_training_single_device.yaml. Loader path is third_party/tt_forge_models/<model_dir>/pytorch/loader.py where <model_dir> is the segment before the first / in the key.
3. model_info.name path: parse as <framework>_<model>_<variant>_<task>_<source>. The framework is always pytorch for this skill. Search loaders with: grep -rln "model=.*\"<model>\"\|ModelName.*<model>" third_party/tt_forge_models/, then for each candidate, read its _get_model_info to confirm the variant/task/source line up. Once the loader is identified, derive the YAML key as <model_dir>/pytorch-<variant>-single_device-training and confirm it exists in the YAML.
4. If neither resolves, abort with: Could not resolve <test-name> to a YAML entry or loader. Provide either the YAML key or the model_info.name.

Step 2 - Read context, do not gate on the YAML reason

The skill is always invoked with an explicit <test-name>, so the user's intent is the source of truth - YAML reasons can be stale (the runner may have masked the real failure behind the canonical unpack_forward_output message, or the entry may have been left in place after a prior fix). Read the YAML entry's reason: for context and include it verbatim in the final report, but do not abort here on a mismatch. The real check happens at Step 3 (CPU triage): if the model successfully runs end-to-end and its forward output is already handled, Step 3 short-circuits the workflow.

Cheap early exit before Step 3: grep -n "def unpack_forward_output(" <loader_path>. If the override already exists, skip to Step 6 with a note that the YAML is stale.

Step 3 - CPU triage

Before running: confirm IRD_LF_CACHE is set, then ask the user if they want to proceed:

Note: The triage script downloads model weights from the LF cache. If IRD_LF_CACHE is not exported, downloads will fail with a confusing network error rather than an actionable message. Check with echo $IRD_LF_CACHE; if empty, run export IRD_LF_CACHE=<link>. Then confirm with the user: "Ready to run the CPU triage for <model_dir>. This will download weights and run a forward pass on CPU - proceed?"

Run the bundled script. It loads the model, calls forward, and prints the output structure plus a final OUTPUT_CLASS: line for grep:

python .claude/skills/triage-unpack-forward-output/scripts/triage_forward_output.py \
    --model-dir <model_dir> [--variant <NAME>] [--batch-size 2] \
    &> /tmp/triage_<model_dir>_<variant>.log

<model_dir> is the segment before the first / in the YAML key (or the multi-segment path for nested loaders like bert/question_answering). --variant is optional - pass the ModelVariant enum name if the loader has one, omit it otherwise. The script resolves the repo root from its own location, so cwd does not have to be the tt-xla root, but source venv/activate must have been run.

Read the log with the Read tool. Never use tail or less - they hide errors behind generic exit codes (this is a hard rule from feedback memory).

If the script errored out before producing structure, the bundled script is simple enough to copy and edit locally (try dtype_override=torch.bfloat16, or seq_len if load_inputs accepts it - see tests/runner/utils/dynamic_loader.py:380-419 for what the test runner passes). Save the local copy to /tmp/triage_<...>.py and run it; do not commit changes to the bundled script for one-off model-specific tweaks.

After Step 3 produces an output class, do these checks before continuing:

• grep -n "<OutputClassName>" third_party/tt_forge_models/training_utils.py. If the class is already registered and the registered attribute lookup yields a tensor, the YAML is stale - skip to Step 6 and note this.
• If the class is in the registry but the registered attribute is itself a tuple/list/dict (the registry entry is wrong, e.g. the historical MgpstrModelOutput.logits 3-tuple), this is the wrong-registry case: in Step 5, plan to delete the bad registry entry and write a per-loader override instead.

Step 4 - Gather model docs

From the loader file, extract the model's HuggingFace id (pretrained_model_name in _VARIANTS), GitHub URL or paper if commented at the top. Use WebFetch on the model card / repo README to identify what training loss conventionally consumes.

Defaults if docs are silent or model fits a known family:

Step 5 - Decide registry vs override, then apply the fix

Universal rule (both paths): return only what the loss depends on. Never include auxiliary outputs (attention weights, hidden-state caches, FPN side branches, raw anchors, debug tensors) just because they are tensors. Including unused tensors inflates the autograd graph and changes what is exercised on TT.

Decision - registry entry (Step 5a) is the default. Override (Step 5b) is the escape hatch.

Use the registry when all of the following hold:

• The forward output is a class instance (a HuggingFace ModelOutput dataclass or any other named dataclass) with a stable class name. Bare list/tuple/dict cannot use the registry.
• The loss-relevant tensor is a single attribute lookup on that class - no concat/stack/slice needed.
• The class name is specific to one model family, OR the chosen attribute is universally the loss target across every model that emits that class. For HF transformer outputs this is almost always true (CausalLMOutputWithPast -> logits everywhere).

Use the per-loader override when any of the following hold:

• Output is a bare list/tuple/dict (no class name to key on).
• Loss-relevant tensors require a structural transform (concat, stack, slice, recursive walk).
• Two models share an output class but legitimately need different loss targets - register one and override the other.
• The output class is custom (defined in the model's own repo, not in transformers) and unlikely to be reused - case-by-case judgment.
• A registry entry already exists for the class but is wrong: the registered attribute is itself a tuple/list/dict (e.g. the historical _register_attr("MgpstrModelOutput", "logits") where .logits is a 3-tuple, not a tensor). In that case, delete the bad registry entry in training_utils.py and write the per-loader override. The deletion is the audit trail - don't leave the wrong entry behind.

Step 5a - Registry entry (preferred path)

Add one line in third_party/tt_forge_models/training_utils.py, alphabetically ordered next to existing _register_attr calls:

_register_attr("<OutputClassName>", "<loss_relevant_attr>")

That's it - no docstring needed; the registry table itself is the audit trail. The class name is what the CPU triage script reported in Step 3 (e.g. DPRContextEncoderOutput). Confirm python -c "from third_party.tt_forge_models.training_utils import _HANDLER_REGISTRY; assert '<OutputClassName>' in _HANDLER_REGISTRY" (or just re-import training_utils cleanly).

Expand the verification scope. A registry entry is global - every model in the codebase that returns this output class is now affected. Find them:

grep -rln "<OutputClassName>" third_party/tt_forge_models/
grep -rln "<OutputClassName>" venv/lib/python3.12/site-packages/transformers/  # for HF classes, to confirm which model families emit it

For every YAML entry whose reason: is the canonical unpack_forward_output message and whose loader returns this same class, verify it as part of Step 6 and update its YAML in Step 8. Do not silently leave stale entries - they will misrepresent the test state.

Step 5b - Per-loader override

Add def unpack_forward_output(self, forward_output) to that model's ModelLoader class. Two non-negotiable rules:

1. Return only what the loss depends on (universal rule above).
2. Docstring with these two sections. The override must carry them - they are the audit trail for why this override exists:
   • Forward output structure - full type/shape signature of what the model returns. Example: list[Tensor] of length 3, shapes [(B, 256, 80, 80), (B, 256, 40, 40), (B, 256, 20, 20)].
   • What is selected and why - which subset is returned and which loss it is the gradient source of. Example: returning the three detection-head tensors concatenated; these are the only outputs consumed by YOLOLoss; auxiliary feature maps are dropped.

Implementation guidance:

• For nested list/tuple/dict structures: import extract_tensors_recursive from third_party.tt_forge_models.tools.utils, walk the loss-relevant subtree, and sum(tensors) - see yolov9/pytorch/loader.py:171-194.
• For a plain tuple where only one element is the loss target: index it.
• Keep the implementation minimal. No extra abstractions, no inline comments beyond the required docstring.

After writing (either path), run pre-commit run --files <changed_files>. If pre-commit is not on PATH, the venv isn't active - run source venv/activate first.

Step 6 - Verify with pytest

Do not pre-edit the YAML. Verification runs the tests with --force-run --runxfail so that the current YAML status is bypassed without an intermediate flip:

• --force-run makes NOT_SUPPORTED_SKIP entries execute (tests/runner/conftest.py:113, tests/runner/test_models.py:115-119, tests/runner/test_utils.py:746-748).
• --runxfail makes the runner's imperative pytest.xfail(reason) (tests/runner/test_utils.py:749-750) report the real outcome instead of XFAIL - covers stale KNOWN_FAILURE_XFAIL entries.

The set of "affected entries" depends on the path you took in Step 5:

• Step 5a (registry): every YAML entry whose loader emits the registered class and currently fails with the unpack_forward_output reason.
• Step 5b (override): just the one entry tied to the loader you edited.

Use --junitxml=<path> to capture structured per-testcase outcomes - the runner records the detailed error (TT_FATAL / OOM / error: lines) at tests/runner/test_utils.py:737 via record_property("error_message", ...), which lands in the XML as <property name="error_message" value="…"/>. Pytest writes any failure traceback into <failure message="…"> separately. The XML is the source of truth in Step 7. Also dump stdout/stderr to a .log companion in case the XML is ambiguous.

Run the affected training tests in a single pytest invocation, plus the inference test for the originally-requested entry:

timeout 600 pytest --force-run --runxfail \
  --junitxml=/tmp/verify_<scope>_training.xml \
  "tests/runner/test_models.py::test_all_models_torch[<entry1>-single_device-training]" \
  "tests/runner/test_models.py::test_all_models_torch[<entry2>-single_device-training]" \
  &> /tmp/verify_<scope>_training.log

timeout 300 pytest --force-run --runxfail \
  --junitxml=/tmp/verify_<scope>_inference.xml \
  "tests/runner/test_models.py::test_all_models_torch[<original_entry>-single_device-inference]" \
  &> /tmp/verify_<scope>_inference.log

Read both XML files with the Read tool. Use the .log companion only if the XML is missing detail (e.g. only a generic "Error code 13" with no error_message). Never write Error code 13 to the yaml.

Step 7 - Classify the training outcome from junit XML

For each <testcase> in the training XML:

• No <failure> / <error> / <skipped> child -> pass, outcome = EXPECTED_PASSING.
• <failure message="…">…</failure> present -> fail. Read both:
  • <failure>'s message attribute and inner text (raw exception + traceback).
  • <property name="error_message" value="…"/> inside <properties> - the runner's pre-extracted detailed error (TT_FATAL / OOM / error: line). This is the canonical one-liner to use as the YAML reason:.
• <skipped type="pytest.xfail" message="…"/> -> unexpected with --runxfail; means the entry is xfailed via a different mechanism. Flag for human review.
• <skipped type="pytest.skip" message="…"/> -> unexpected with --force-run; means a different skip path triggered. Flag for human review.

Pick the matching bringup_status from the failure text:

If the XML's error_message is missing or only contains a generic exit-code message, fall back to the .log companion file: grep -nE "TT_FATAL|TT_THROW|error:" /tmp/verify_<scope>_training.log and use the first matching line for context.

Step 8 - Update the YAML

Edit tests/runner/test_config/torch/test_config_training_single_device.yaml. The YAML was not pre-edited in Step 6, so this is a single write per affected entry. Update every training entry covered by Step 6 - not just the original one. The goal: get tests out of NOT_SUPPORTED_SKIP so they actually run in nightly/weekly. Pick by outcome:

• Pass:

  <key>:
    status: EXPECTED_PASSING
  

  Drop bringup_status and reason entirely.
• Fail with INCORRECT_RESULT (PCC failure): the test is functionally running, just inaccurate. Run it as EXPECTED_PASSING with PCC checks disabled:

  <key>:
    status: EXPECTED_PASSING
    assert_pcc: false
  

  Drop bringup_status and reason.
• Any other failure (FAILED_TTMLIR_COMPILATION, FAILED_RUNTIME, or even still FAILED_FE_COMPILATION for an unrelated env/import reason like _has_torch_function docstring, is_torch_fx_available, Boolean value of Tensor … ambiguous, ModuleNotFoundError): mark as KNOWN_FAILURE_XFAIL so nightly/weekly invokes the test and reports XFAIL with the documented reason instead of silently skipping:

  <key>:
    status: KNOWN_FAILURE_XFAIL
    bringup_status: <FAILED_FE_COMPILATION | FAILED_TTMLIR_COMPILATION | FAILED_RUNTIME>
    reason: "<one-line excerpt of the real error, ~120 chars>"
  

  Do not add a markers: field unless one was already present for an unrelated reason. Never leave a triaged entry at NOT_SUPPORTED_SKIP - the runner skips those (tests/runner/test_models.py:118) and the failure is invisible in nightly reports.

The inference log is informational only. Do not modify the inference YAML entry. Flag any inference regression to the user in the final report so they can investigate separately.

Step 9 - Final report

Print a single concise summary to the user:

1. Resolved test -> loader path + YAML key
2. Forward output structure observed on CPU (one or two lines)
3. Fix path taken: registry entry (which class -> which attr, which other YAML entries it covered) or per-loader override (file, method signature, what it returns). State explicitly which path you chose and why the other was not used.
4. Pytest outcome for every affected training entry (pass/fail with the real error from the junit XML's error_message if fail), plus the original inference entry.
5. YAML diff (the lines that changed across all affected entries).
6. Anything that needs human review (e.g. inference regression, ambiguous loss target, registry entry that affected an unexpected number of tests, extract_tensors_recursive returned an unexpected number of tensors).

Hard rules (do not violate)

• Invoked with one test at a time. If <test-name> is missing or matches multiple entries, abort and ask.
• Only act when Step 3's CPU triage shows the model returns a class/structure that is not already correctly handled by unpack_forward_output. The YAML reason: is informational - do not gate on it (it may be stale). If Step 3 shows the output is already covered by an existing handler that yields a tensor, the YAML is the only thing wrong: skip to Step 6 and update the YAML.
• Prefer the registry path (Step 5a) over per-loader override (Step 5b). Use the override only when the decision criteria in Step 5 require it.
• Per-loader override must return only loss-relevant tensors and must carry the two-section docstring defined in Step 5b (forward output structure, what is selected and why).
• Whichever path is taken, the fix must return only loss-relevant tensors - never aux outputs (attentions, hidden states, FPN side branches).
• A registry entry has global scope. Re-verify and update YAML for every affected training entry, not just the originally requested one.
• Verification uses pytest --force-run --runxfail --junitxml=<path>. Do not pre-edit the YAML before pytest. Read the XML for outcomes; classify failures from the <failure> element and the error_message property.
• Touch only the training YAML entries. Inference is informational.
• Do not commit or push. Leave changes staged-but-uncommitted unless the user asks otherwise.