Add an event-log witness-tracing analysis rule

[fzyzcjy]

Add the witness-tracing rule for the event analyzer: it follows witness ids
through the replayed event log to verify they are propagated correctly across the
training pipeline, with unit tests.

• miles/utils/event_analyzer/rules/witness.py and tests.

[gemini-code-assist]

Code Review

This pull request introduces a new witness event analyzer rule and its associated tests to detect witness data mismatches and missing snapshots during training. The feedback focuses on optimizing performance by eliminating nested loop bottlenecks in mismatch detection, adding robust shape and length validation checks for advantage and witness ID lists, and utilizing in-place set operations to improve efficiency.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Optimize performance by eliminating two nested loop bottlenecks in _find_mismatches:

1. Group latest_attempt_witness_events by (rollout_id, cell_index) beforehand to reduce the lookup from $O(W)$ to $O(1)$.
2. Precompute the zero-advantage excused IDs for all unique rollout_ids in a single sorted pass to avoid the $O(N^2)$ complexity of calling _zero_adv_excused_ids_at repeatedly.

    latest_attempt_witness_events = _filter_to_latest_attempt(
        all_witness_events, group_key=lambda e: (e.rollout_id, e.source.cell_index)
    )
    witness_events_by_key = defaultdict(list)
    for e in latest_attempt_witness_events:
        witness_events_by_key[(e.rollout_id, e.source.cell_index)].append(e)

    # Precompute zero-advantage excused IDs for each rollout_id in a single pass to avoid O(N^2) complexity
    zero_adv_excused_ids_by_rollout_cache: dict[int, set[int]] = {}
    excused: set[int] = set()
    all_rids = sorted(set(zero_adv_witness_ids_by_rollout) | set(allocated_witness_ids_by_rollout))
    rid_idx = 0
    n_rids = len(all_rids)
    unique_rollout_ids = sorted({step_event.rollout_id for step_event in all_step_events})
    for r_id in unique_rollout_ids:
        while rid_idx < n_rids and all_rids[rid_idx] <= r_id:
            rid = all_rids[rid_idx]
            excused -= allocated_witness_ids_by_rollout.get(rid, set())
            excused |= zero_adv_witness_ids_by_rollout.get(rid, set())
            rid_idx += 1
        zero_adv_excused_ids_by_rollout_cache[r_id] = set(excused)

    for step_event in all_step_events:
        rollout_id = step_event.rollout_id

        for cell_index, cell_outcome in step_event.cell_outcomes.items():
            if cell_outcome == "error":
                continue
            if not all(r == TrainStepOutcome.NORMAL for r in cell_outcome):
                continue

            witness_events_of_cell = witness_events_by_key.get((rollout_id, cell_index), [])

            if not witness_events_of_cell:
                yield WitnessMissingSnapshotIssue(
                    rollout_id=rollout_id,
                    cell_index=cell_index,
                    description=f"Cell {cell_index} reported NORMAL for rollout {rollout_id} but no WitnessSnapshotParamEvent was found",
                )
                continue

            zero_adv_excused_ids = zero_adv_excused_ids_by_rollout_cache.get(rollout_id, set())

[gemini-code-assist]

When processing lists of sequence-level or token-level tensors (such as advantages and witness IDs), validate that the list lengths match up front, and perform per-sample shape checks to prevent silent mismatches or broadcast failures.

Suggested change

	for adv_tokens, wid_tokens in zip(event.advantages, event.witness_ids, strict=True):
	if adv_tokens and all(v == 0.0 for v in adv_tokens):
	result[event.rollout_id].add(wid_tokens[0])
	if len(event.advantages) != len(event.witness_ids):
	raise ValueError("Length mismatch between advantages and witness_ids")
	for adv_tokens, wid_tokens in zip(event.advantages, event.witness_ids, strict=True):
	if len(adv_tokens) != len(wid_tokens):
	raise ValueError("Shape mismatch in sequence-level tokens")
	if adv_tokens and wid_tokens and all(v == 0.0 for v in adv_tokens):
	result[event.rollout_id].add(wid_tokens[0])

References

1. When processing lists of sequence-level or token-level tensors (such as advantages, student log probabilities, and teacher log probabilities) in RL or distillation pipelines, validate that the list lengths match up front, and perform per-sample shape checks to prevent silent mismatches or broadcast failures.

[gemini-code-assist]

Optimize set union by using the in-place operator |= instead of creating a new set object with | in each iteration of the loop.

Suggested change

	for rollout_id in sorted(allocated_witness_ids_by_rollout.keys()):
	running = running | allocated_witness_ids_by_rollout[rollout_id]
	ans[rollout_id] = set(running)
	for rollout_id in sorted(allocated_witness_ids_by_rollout.keys()):
	running |= allocated_witness_ids_by_rollout[rollout_id]
	ans[rollout_id] = set(running)

[Shi-Dong]

Please fix.

[gemini-code-assist]

Remove the _zero_adv_excused_ids_at helper function as it is no longer needed after precomputing the excused IDs in _find_mismatches.

[Shi-Dong]

LGTM with nits.

[Shi-Dong]

Do you think it'd be safer to assert that every entry in wid_tokens is the same here?