[feat] SID: add SidRqkmeans model (FAISS-trained residual K-Means)

[WhiteSwan1]

Second of three PRs splitting the Semantic-ID models onto the shared base from #538. Adds the concrete RQ-KMeans backend on top of ResidualQuantizer / BaseSidModel; RQ-VAE follows in PR3.

• tzrec/modules/sid/kmeans.py: KMeansLayer centroid container + recon_diagnostics.
• tzrec/modules/sid/residual_kmeans_quantizer.py: ResidualKMeansQuantizer (FAISS-trained, FX-traceable forward, non-uniform per-layer codebooks).
• tzrec/models/sid_rqkmeans.py: SidRqkmeans(BaseSidModel) - gradient -free; reservoir-samples embeddings during the train loop and fits FAISS once in on_train_end.
• tzrec/models/model.py: BaseModel.on_train_end() no-op lifecycle hook.
• tzrec/main.py: invoke on_train_end after the train loop and force the tail checkpoint so post-hook state is persisted.
• protos: SidRqkmeans message + ModelConfig registration (601; 600 is reserved for SidRqvae in PR3).
• tests: kmeans_test, ResidualKMeansQuantizerTest, sid_rqkmeans_test.

[github-actions]

Review summary — SidRqkmeans (FAISS residual K-Means)

Solid, carefully engineered PR. The Vitter Algorithm R reservoir, the int-status-flag DDP fit/broadcast (sidestepping NCCL bool quirks), the cached _initialized mirror to avoid per-batch .item() syncs, and the mid-fit-checkpoint poison guard are all thoughtful and well-documented. Two concrete issues + a few test/doc notes below.

Should-fix (inline):

• FAISS never runs on GPU — gpu=torch.cuda.current_device() passes a device index, but faiss treats gpu as a GPU count where 0 (the common case, incl. the rank0-only DDP fit) is falsy → silent CPU fallback. The codebook fit is typically the largest one-shot cost, so this is a real perf regression.
• gather_object on the NCCL group — should use the existing dist_util.get_dist_object_pg() gloo helper, as odps_dataset.py does for the same pattern.

Test gaps worth closing:

• The reservoir Phase-2 steady-state replacement (accept prob cap/(n_seen+j+1), float64 exactness, slot collisions — the subtlest code in the PR) is never validated for correctness; test_reservoir_caps_memory only checks counts/shape, which pass regardless of replacement logic.
• normalize_residuals=True is never exercised through train_offline — note its F.normalize is a second normalize site independent of _residual_pass, so the two can silently diverge.
• The eval-metric path (init_metric/update_metric, mse/rel_loss, the eval-only quantized/input_embedding outputs) and the inference-mode (set_is_inference) codes-only contract are untested.
• test_post_fit_checkpoint_round_trips asserts only codes.abs().sum() > 0; comparing against the source model's codes (torch.testing.assert_close) would actually verify the round-trip rather than just "not uninitialized".

Minor (docs/consistency):

• on_train_end Returns: doc says False on empty reservoir, but the DDP path never returns False (it raises via the fail-fast status broadcast) — worth clarifying the two paths differ.
• train_offline docstring ("moves only FAISS's subsampled working set to the GPU") understates the chunked index.search over all N rows; and the train_sample_size proto comment says K * max_points_per_centroid but the code uses max(K) * ... for non-uniform codebooks.

[github-actions]

Review summary

Strong, carefully-engineered PR. The hard parts are handled well: FX-traceable read-only predict, the cached _initialized bool to avoid a per-batch GPU→CPU .item() sync, Phase-2 reservoir sampling that copies only accepted rows to host, the mid-fit-checkpoint poison rejection in _load_from_state_dict, and a deadlock-free DDP on_train_end (status flag broadcast before the centroid broadcast, int64 not bool for NCCL). Test coverage is unusually thorough (reservoir Phase-2 correctness, non-uniform codebooks, normalize path, exact checkpoint round-trip, 2-rank DDP happy path). I ran five focused review passes; most areas came back clean.

Three items worth a look (posted inline):

1. gpu=1 selects all GPUs, not rank0's device (residual_kmeans_quantizer.py) — verified against faiss source: gpu=1 collapses to True (since 1 == True) → get_num_gpus() → index_cpu_to_all_gpus(ngpu=ALL). With all GPUs visible per rank, the rank0-only fit allocates faiss temp memory across every rank's GPU. Results stay correct, but it contradicts the comment and risks OOM on tight runs. This is in the line the latest commit just changed, and CI (faiss-cpu, 0 GPUs) can't exercise it.

2. The DDP coordinated-failure path is untested (sid_rqkmeans.py) — the status-flag-before-broadcast logic exists specifically to avoid a deadlock; a regression would hang rather than fail, with no test to catch it.

3. The is_ckpt_after_train tail checkpoint is untested (main.py) — it's the only thing persisting the codebook for this FAISS-only model.

Minor (not posted inline):

• _init_reservoir docstring says the per-rank cap targets default_fit_sample_size(), but the code targets train_sample_size when it's set (>0) and only falls back to the default at 0 — worth a one-line tweak.
• The faiss end-to-end tests (test_on_train_end_runs_faiss, non-uniform, normalize, eval-contract) don't seed RNG and assert only "codes in range / centroids non-zero"; seeding + one reconstruction-error threshold would make them deterministic and able to catch a broken residual walk.
• A misspelled faiss_kmeans_kwargs key (e.g. n_iter) only fails at fit time — i.e. after a full training run. Validating keys at construction would fail fast.

Nice work overall — the engineering rigor here is clearly high.

[github-actions]

Code Review Summary

Reviewed by parallel subagents (code quality, performance, test coverage, documentation accuracy, security), with findings cross-checked before posting.

Overall this is carefully written code. The fail-fast guards (reservoir cap vs. largest K at construction), the mid-fit checkpoint poison detection, the strictly-typed FAISS kwargs proto, the bounded reservoir with float64 stream counters, and the exact-round-trip checkpoint tests all show real attention to edge cases. Test coverage of the prediction contracts, non-uniform codebooks, and checkpoint poisoning is genuinely strong.

Main concern (inline on tzrec/main.py)

Four of five reviewers independently converged on the same issue: the "periodic checkpointing disabled" assumption is documented in three places but enforced nowhere. save_checkpoints_steps defaults to 1000, maybe_save(final=True) is still subject to the per-step dedupe, and a periodic save landing on the final step silently discards the fitted codebook — the surviving checkpoint loads cleanly and emits all-zero SIDs. This is the one silent-data-loss path in the PR and worth closing before merge (details + fix options inline).

Lower-priority items (no inline comment)

• import tzrec now hard-depends on faiss at import time: residual_kmeans_quantizer.py imports faiss at module level and auto_import eagerly loads all model modules. Since faiss is in all requirements files this works, but it makes the eleven skipTest("faiss not installed") guards in the new tests unreachable (collection would fail first), and faiss.contrib.torch_utils monkey-patches faiss process-wide as an import side effect. Moving the import inside train_offline would decouple this.
• x_hat is exposed even with normalize_residuals=True, where the inline comment itself notes reconstruction metrics are not meaningful — mse/rel_loss then compare across scales with no warning. Consider gating or warning.
• Reservoir subsampling uses the unseeded global RNG, so once the corpus exceeds the cap, two runs with identical config (including faiss seed) produce different codebooks/SIDs. A config-seeded torch.Generator would make SID generation reproducible.
• train_offline's out buffer exists only for the verbose log, and the per-layer diagnostics materialize several extra (N, D) temporaries — fine at the default cap, but a large train_sample_size multiplies peak host memory ~3-4x right at the end of the run. Folding scalar diagnostics into the existing chunk loop would fix it.
• Integration test asserts only that eval_result.txt exists — an unfitted-restore regression produces NaN metrics but still passes. Parsing the file and asserting mse/rel_loss are finite would catch that whole class (including the dedupe issue above).
• No user-facing docs: no docs/source/models/ page or README model-table row for the new sid_rqkmeans config. Fine if planned for PR3, but the save_checkpoints=0 requirement especially needs documenting somewhere users will see it.

🤖 Generated with Claude Code

[github-actions]

Code Review Summary

Reviewed at head a9a889c by parallel subagents (code quality, performance, test coverage, docs accuracy, multi-process safety), findings cross-checked before posting. Prior-round items were verified first: the DDP coordination concerns are resolved by removal (single-process guard in __init__, raised on all ranks before any collective), the faiss gpu kwarg is stripped, tests are now seeded with real assertions, the reservoir-cap and faiss-kwargs validation landed, and the checkpoint-dedupe risk is explicitly documented as a TODO. The fail-fast posture of this commit (raise-not-assert with the python -O rationale, typed faiss kwargs, width guards) is genuinely good work.

Inline comments (5)

1. CUDA-guard message recommends the workaround this commit proved unreliable — the error says CUDA_VISIBLE_DEVICES="", while the integration-test fix in the same commit switched to "-1" because "" doesn't hide devices on the GPU CI runtime. One-line fix.
2. NaN/Inf embeddings defeat the poison guard (residual_kmeans_quantizer.py) — a non-finite row from external data becomes NaN centroids, and the all-zero poison check (abs().sum() == 0) passes a NaN codebook as healthy → silent degenerate SIDs. One torch.isfinite check in train_offline closes it. Also: the adjacent comment mis-states faiss (it throws for N < K; warn-only is N < K * min_points_per_centroid).
3. Empty-vs-small reservoir inconsistency (on_train_end) — empty warns and exits 0 with an unfitted checkpoint persisted; 0 < N < max(K) hard-fails after the whole pass. The warn path is the same silent zero-SID artifact the dedupe TODO describes; suggest raising in both cases.
4. The new fail-fast validations have no negative tests — all four BaseSidModel ValueError paths and both train_offline RuntimeError guards are uncovered; a revert passes the suite. ~4 lines each with the existing helpers.
5. Integration test docstring overstates coverage — it claims finite-metric verification but asserts only file existence; parsing eval_result.txt for finite mse/rel_loss would also give the dedupe TODO an end-to-end regression net.

Minor (no inline)

• train_offline still allocates and accumulates the (N, D) out buffer when verbose=False (the commit gated the x0 clone but not out); in the default normalize_residuals=False path the logged loss reduces to mean(x**2), making out avoidable entirely.
• RelativeL1.update has no shape check — broadcasting silently corrupts the metric; torchmetrics convention is _check_same_shape at the top of update. The float64-accumulation fix also has no regression test.
• BaseSidModel class docstring lists mse/unique_sid_ratio but omits rel_loss (the init_metric docstring is correct).
• Reservoir subsampling still uses the unseeded global RNG, so identical configs (incl. faiss seed) give different codebooks — noted in a prior round; fine if accepted, but the proto seed field implies determinism it can't deliver.
• Docs page / README model-table row still absent — assuming this is deferred to the PR3 wrap-up; the save_checkpoints_steps: 0 and CPU-only/single-process constraints especially need a user-visible home.

🤖 Generated with Claude Code

[WhiteSwan1]

Thanks for both github-actions passes — went through them all. Dispositions:

Fixed:

• proto wording + codebook≥1/input_dim≥1 validation;
• removed the duplicated relative-L1 helper (the metric is now the single source);
• float64/long RelativeL1 accumulators;
• assert→raise for the -O-stripped data-corruption guards;
• mock config save_checkpoints=0;
• corrected the faiss-N<K comment;
• guard message now says CUDA_VISIBLE_DEVICES="-1";
• added negative tests for every fail-fast path and strengthened the integration test to assert finite reconstruction metrics.

Declined / deferred (with reason):

• __init__ CPU/world_size guards stay in __init__ — v1 is intentionally CPU-only end to end (training and inference), so failing fast at every entry point is the intended contract. Because KMeans inference is so simple and only involves nearest-neighbor search, it's unlikely we'll need GPU inference down the road.
• Checkpoint-dedupe enforcement — kept the save_checkpoints=0 convention as agreed; left a TODO to harden it in a follow-up.
• ReservoirSampler (B,1) width guard — the embedding width is never 1 in practice, so reverted.
• Empty-vs-small-reservoir consistency and a NaN/Inf input check — reasonable, but behavior/logic changes better suited to a follow-up.

Also fixed the CPU-only end-to-end test that was failing on the GPU CI image: it now runs on the CPU CI job and skips where CUDA is present.

[WhiteSwan1]

Thanks @tiankongdeguiji — all addressed on latest (5f5af01):

Round 1

1. input_embedding in predictions — removed; predict exposes only the reconstruction x_hat (eval, once fitted), and update_metric re-extracts the target from the batch.
2. Integration test — moved to tzrec/tests/sid_integration_test.py (with a mock config); it also asserts the post-fit eval reports finite metrics.
3. init_metric/update_metric → BaseSidModel; rel_loss is now a torchmetric — a symmetric RelativeL1 rather than MeanAbsolutePercentageError (MAPE's |t-p|/|t| is asymmetric and unbounded as target→0; the symmetric form matches OpenOneRec's calc_loss). Happy to switch to the builtin if you prefer.
4. DDP — removed; CPU-only, single-process, with a world_size>1 fail-fast in __init__.
5. torch.cdist(x, centroids).argmin(-1) — adopted; _squared_euclidean_distance deleted.

Round 2

1. File renamed to residual_kmeans_quantizer.py.
2. copy=True — removed; train_offline now explicitly consumes its input (documented), and the caller decides whether to copy.
3. Added a QuantizeLayer base class (tzrec/modules/sid/quantize_layer.py); KMeansQuantizeLayer subclasses it and the PR3 VectorQuantizeLayer will too.
4. Replaced the Struct + _coerce_proto_numbers with a typed FaissKmeansConfig proto message (strictly-typed faiss kwargs).
5. Dropped the force param; SID sets save_checkpoints_steps = save_checkpoints_epochs = 0.