more

Author: albertz

users/zeyer/experiments/exp2025_07_07_in_grads/jobs/chunk_segmentation.py

@@ -377,9 +377,17 @@ class ChunkSegmentationFromModelBatchedJob(Job):
     in bf16 it may differ by matmul non-determinism (set ``model_dtype`` in the model config).
     """
 
-    # Bumped to 2 when sort_by_length / batched_logprobs were added: the flags are part of the hash
-    # (no __sis_hash_exclude__), so every value combination gets its own hash.
+    # Bumped to 2 when sort_by_length / batched_logprobs were added:
+    # those flags are part of the hash, so every value combination gets its own hash.
     __sis_version__ = 2
+    # Algorithm variants added later.
+    # The excluded value for each is the one that REPRODUCES the behavior from before the param existed
+    # (no beam / no bias / summed word scores),
+    # so a job left at the defaults hashes like a pre-variant job and the finished sweeps stay reused;
+    # a non-default value still changes the hash.
+    # Polarity matters:
+    # excluding the *new* value instead would make a variant job collide with the old hash and not rerun.
+    __sis_hash_exclude__ = {"word_start_beam": None, "exit_bias": 0.0, "length_norm": False}
 
     def __init__(
         self,
@@ -395,6 +403,9 @@ def __init__(
         max_batch_size: int = 8,
         sort_by_length: bool = True,
         batched_logprobs: bool = True,
+        word_start_beam: Optional[float] = None,
+        exit_bias: float = 0.0,
+        length_norm: bool = False,
     ):
         super().__init__()
         self.dataset_dir = dataset_dir
@@ -411,6 +422,19 @@ def __init__(
         # instead of one per word (the DP's real bottleneck). Both default on (verified equivalent).
         self.sort_by_length = sort_by_length
         self.batched_logprobs = batched_logprobs
+        # word_start_beam: keep every previous-chunk exit within this log-prob margin of the best,
+        # and start from the earliest survivor.
+        # None = plain argmax (the original heuristic);
+        # a wider beam prunes less and approaches word_start_heuristic=False (exact).
+        # One knob interpolating the two, since the argmax end measurably beats the exact end.
+        # exit_bias: added to EVERY exit log-prob,
+        # unlike empty_exit_penalty which only applies to a chunk exiting with zero words --
+        # a global words-per-chunk knob (cf. word insertion penalty).
+        # length_norm: score a word by its per-token MEAN log-prob instead of the sum,
+        # so emitting a multi-token word is not systematically dearer than the single exit token.
+        self.word_start_beam = word_start_beam
+        self.exit_bias = exit_bias
+        self.length_norm = length_norm
 
         self.rqmt = {"time": 40, "cpu": 2, "gpu": 1, "mem": 125}
         self.out_hdf = self.output_path("out.hdf")
@@ -537,9 +561,15 @@ def _chunk_bounds(audio, samplerate):
                     if cci == 0 or not self.word_start_heuristic:
                         prev_idx, cws = 0, 0
                     else:
-                        prev_idx = int(
-                            torch.stack([n.accum_exit_log_prob for n in s["array"][cci - 1]]).argmax().item()
-                        )
+                        exits = torch.stack([n.accum_exit_log_prob for n in s["array"][cci - 1]])
+                        if self.word_start_beam is None:
+                            prev_idx = int(exits.argmax().item())
+                        else:
+                            # Beam prune:
+                            # keep every exit within the margin of the best, start from the earliest survivor.
+                            # Wider beam -> prunes less -> approaches the exact DP.
+                            keep = (exits >= exits.max() - self.word_start_beam).nonzero()
+                            prev_idx = int(keep.min().item())
                         cws = s["array"][cci - 1][prev_idx].word_idx
                     cwe = len(s["words"])
                     if cws >= cwe:
@@ -557,9 +587,10 @@ def _chunk_bounds(audio, samplerate):
                 if not active:
                     break
 
-                # One flushed line per chunk step, so the log keeps updating and a slow run is
-                # distinguishable from a real hang: small chunk sizes mean ~1000+ chunks per seq, and
-                # the DP is otherwise silent for a whole group (only group start / seq end print).
+                # One flushed line per chunk step, so the log keeps updating
+                # and a slow run is distinguishable from a real hang:
+                # small chunk sizes mean ~1000+ chunks per seq,
+                # and the DP is otherwise silent for a whole group (only group start / seq end print).
                 print(
                     f"** chunk step: {len(active)} active seqs,"
                     f" chunk {[st[i]['cci'] for i in active]} of {[len(st[i]['cse']) for i in active]},"
@@ -606,9 +637,13 @@ def _chunk_bounds(audio, samplerate):
                                 0.0,
                             )
                             wlp = wlp.sum()
+                            if self.length_norm:
+                                wlp = wlp / max(1, int(t1 - t0))
                         else:
                             wlp = None
                         exit_lp = log_probs[0, 0, model.assistant_end_token_id]
+                        if self.exit_bias:
+                            exit_lp = exit_lp + self.exit_bias
                         if w == 0:
                             exit_lp = exit_lp + self.empty_exit_penalty
                         prev_left, prev_below = None, None

users/zeyer/experiments/exp2026_05_22_chunk_align.py

@@ -221,6 +221,11 @@ def py():
         # the -5 default fights that): does removing the penalty recover acc 0.07 / 0.02?
         (1.0, 0.5, [0.0, -5.0]),
         (0.5, 0.0, [0.0, -5.0]),
+        # optimal overlap at cs1 once eep is set right (eep=0): ov0.5 gave 0.5530, so map the rest.
+        # eep=-5 for cs1-ov0 is the plain cs1-ov0 job above (0.6318).
+        (1.0, 0.0, [0.0]),
+        (1.0, 0.25, [0.0]),
+        (1.0, 0.75, [0.0]),
     ]:
         for _eep in _eeps:
             _seg_e = ChunkSegmentationFromModelBatchedJob(
@@ -280,6 +285,91 @@ def py():
             reg(f"{_x_name}-accuracy.txt", _metric_x.out_accuracy)
             reg(f"{_x_name}-chunk_idx_mae.txt", _metric_x.out_chunk_idx_mae)
 
+    # --- algorithm variants (batched job only) ---
+
+    # 1) word_start_beam: one knob interpolating the argmax heuristic (beam=None, the old default)
+    # and the exact DP (word_start_heuristic=False).
+    # cs1-ov0 is where the two endpoints differ most (0.6318 argmax vs 0.2473 exact), so map the curve there:
+    # wider beam = less pruning = closer to exact, which we expect to get WORSE
+    # (search-error effect, cf. the NMT beam-search curse).
+    for _beam in [0.0, 1.0, 2.0, 5.0, 10.0, 20.0]:
+        _seg_bm = ChunkSegmentationFromModelBatchedJob(
+            dataset_dir=dl_ds_buckeye.out_hub_cache_dir,
+            dataset_key="val",
+            model_config=_cfg_hp,
+            chunk_size_secs=1.0,
+            chunk_overlap_secs=0.0,
+            max_batch_size=8,
+            word_start_beam=_beam,
+        )
+        _bm_name = f"chunk-align/phi4mm-buckeye-val-cs1-ov0-beam{_beam:g}"
+        _seg_bm.add_alias(_bm_name)
+        reg(f"{_bm_name}.hdf", _seg_bm.out_hdf)
+        _m_bm = CalcChunkAssignmentMetricsJob(
+            chunk_seg_hdf=_seg_bm.out_hdf,
+            dataset_dir=dl_ds_buckeye.out_hub_cache_dir,
+            dataset_key="val",
+            dataset_offset_factors=_DATASET_OFFSET_FACTORS["buckeye"],
+        )
+        _m_bm.add_alias(f"{_bm_name}-metric")
+        reg(f"{_bm_name}-accuracy.txt", _m_bm.out_accuracy)
+        reg(f"{_bm_name}-chunk_idx_mae.txt", _m_bm.out_chunk_idx_mae)
+
+    # 2) exit_bias: a global words-per-chunk knob (added to EVERY exit),
+    # vs empty_exit_penalty which only hits chunks exiting with zero words,
+    # and which we showed is a pure heuristic artifact.
+    # Swept with eep=0 so the bias is the only exit knob.
+    # bias=0 is the existing cs2-ov0-eep0 (0.8535).
+    for _bias in [-2.0, -1.0, 1.0, 2.0]:
+        _seg_eb = ChunkSegmentationFromModelBatchedJob(
+            dataset_dir=dl_ds_buckeye.out_hub_cache_dir,
+            dataset_key="val",
+            model_config=_cfg_hp,
+            chunk_size_secs=2.0,
+            chunk_overlap_secs=0.0,
+            empty_exit_penalty=0.0,
+            max_batch_size=8,
+            exit_bias=_bias,
+        )
+        _eb_name = f"chunk-align/phi4mm-buckeye-val-cs2-ov0-eep0-bias{_bias:g}"
+        _seg_eb.add_alias(_eb_name)
+        reg(f"{_eb_name}.hdf", _seg_eb.out_hdf)
+        _m_eb = CalcChunkAssignmentMetricsJob(
+            chunk_seg_hdf=_seg_eb.out_hdf,
+            dataset_dir=dl_ds_buckeye.out_hub_cache_dir,
+            dataset_key="val",
+            dataset_offset_factors=_DATASET_OFFSET_FACTORS["buckeye"],
+        )
+        _m_eb.add_alias(f"{_eb_name}-metric")
+        reg(f"{_eb_name}-accuracy.txt", _m_eb.out_accuracy)
+        reg(f"{_eb_name}-chunk_idx_mae.txt", _m_eb.out_chunk_idx_mae)
+
+    # 3) length_norm: score a word by its per-token MEAN log-prob instead of the sum,
+    # so a multi-token word is not systematically dearer to emit than the single exit token.
+    # Compare against the plain cs*-ov* jobs (cs2-ov0 0.8747, cs10-ov2.5 0.9856, cs1-ov0 0.6318).
+    for _cs, _ov in [(2.0, 0.0), (10.0, 2.5), (1.0, 0.0)]:
+        _seg_ln = ChunkSegmentationFromModelBatchedJob(
+            dataset_dir=dl_ds_buckeye.out_hub_cache_dir,
+            dataset_key="val",
+            model_config=_cfg_hp,
+            chunk_size_secs=_cs,
+            chunk_overlap_secs=_ov,
+            max_batch_size=8,
+            length_norm=True,
+        )
+        _ln_name = f"chunk-align/phi4mm-buckeye-val-cs{_cs:.0f}-ov{_ov:g}-lnorm"
+        _seg_ln.add_alias(_ln_name)
+        reg(f"{_ln_name}.hdf", _seg_ln.out_hdf)
+        _m_ln = CalcChunkAssignmentMetricsJob(
+            chunk_seg_hdf=_seg_ln.out_hdf,
+            dataset_dir=dl_ds_buckeye.out_hub_cache_dir,
+            dataset_key="val",
+            dataset_offset_factors=_DATASET_OFFSET_FACTORS["buckeye"],
+        )
+        _m_ln.add_alias(f"{_ln_name}-metric")
+        reg(f"{_ln_name}-accuracy.txt", _m_ln.out_accuracy)
+        reg(f"{_ln_name}-chunk_idx_mae.txt", _m_ln.out_chunk_idx_mae)
+
     # fp32 batched (default fast path) at cs30, to check the fast path (esp. batched_logprobs) is
     # bit-exact vs the fp32 single-seq reference below. The bf16 sweep diverges more for small
     # chunks, so this isolates real logic differences from bf16 numerical noise.