[feat] SID generation — RQ-VAE and RQ-KMeans models (FAISS-trained)

[WhiteSwan1]

Summary

Introduces two new SID-generation models to tzrec plus the supporting tzrec.modules.sid_generation module stack:

• SidRqvae — end-to-end RQ-VAE (Encoder + Residual VQ + Decoder), trainable via STE or Gumbel-Softmax, with optional CLIP contrastive learning for paired multimodal inputs.
• SidRqkmeans — encoder-free residual K-Means, trained in one shot by FAISS at the end of the train_eval loop ( gradient-free).

Both produce N-layer semantic IDs (code_0, …, code_{n_layers-1}) that downstream generative-retrieval models can consume.

What's in scope

• Models (tzrec/models/):

  • sid_rqvae.py + sid_rqvae_test.py — SidRqvae(BaseModel), no-clip + mixed clip+recon paths.
  • sid_rqkmeans.py + sid_rqkmeans_test.py — SidRqkmeans(BaseModel), FAISS-only training.
  • _sid_helpers.py — shared parse_int_list / parse_float_list for SID proto fields.
  • model.py — adds BaseModel.on_train_end() no-op lifecycle hook (overridden by SidRqkmeans for the FAISS fit).

• Modules (tzrec/modules/sid_generation/):

  • rqvae.py — RQVAE (encoder/decoder MLP + quantizer), supports use_clip mixed-mode.
  • residual_quantized.py — gradient-trained multi-layer residual VQ used by RQ-VAE.
  • residual_kmeans.py — ResidualKMeans + RQKMeans wrapper (FAISS-trained centroid store).
  • kmeans.py — KMeansLayer centroid container + KMeans / KMeans++ utilities.
  • vector_quantize.py — single VQ layer with Sinkhorn uniform assignment, STE / Gumbel-Softmax forward.
  • clip_loss.py — CLIPLoss and MaskedCLIPLoss (gradient-preserving GatherLayer for DDP all-gather, three-way contrastive loss).
  • types.py — QuantizeForwardMode, QuantizeOutput, ResidualQuantizedOutput.

• Protos (tzrec/protos/models/sid_model.proto + tzrec/protos/model.proto):

  • SidRqvae with sub-messages SinkhornConfig and ClipConfig.
  • SidRqkmeans (codebook + normalize_residuals + faiss_kmeans_kwargs Struct).
  • Both registered under ModelConfig.

• Wiring (tzrec/main.py):

  • Two-line addition: _model.on_train_end() is called once after train_and_evaluate exits the loop. No-op for every model except SidRqkmeans, which uses the hook to fit FAISS over the buffered embeddings.

Test plan

• Unit tests — python -m unittest tzrec.models.sid_rqkmeans_test tzrec.models.sid_rqvae_test → 11/11 pass (5 RQKMeans, 6 RQVAE; covers train / eval / inference modes, CLIP all-recon and all-clip edge cases, backward pass).
• Lifecycle smoke — python ft_scripts/lifecycle_smoke.py drives construct → init_loss → init_metric → train (predict/loss/backward/update_train_metric) → on_train_end → eval (predict/loss/update_metric/compute_metric) → set_is_inference(True) → infer, for SidRqkmeans, SidRqvae no-clip, and SidRqvae CLIP. All three pass.
• Backward-compat — existing model configs in master are unaffected; the two new entries are additions to ModelConfig's oneof.

[CLAassistant]

All committers have signed the CLA.

[github-actions]

Critical — resume-then-infer silently returns zero codes.

_is_initialized is persistent=False, so on checkpoint resume the buffer is reset to False even though centroids was restored. ResidualKMeans.forward then takes the else branch and returns torch.zeros(...) for every layer (residual_kmeans.py:135-138) — no exception, no warning. The comment above accurately predicts this scenario ("would silently return dummy zero codes") but the code does not raise either; it just emits zeros.

There is no set_extra_state/on_load_state_dict hook that re-flips this flag after a load, so users running evaluate/predict against a saved RQKMeans checkpoint get garbage SIDs.

Suggested fix: derive initialization from centroids.any() at use sites, or raise in forward/get_codes when not is_initialized rather than silently returning zeros. Otherwise resume-then-eval is broken.

[github-actions]

ResidualQuantized.__init__ asserts commitment_loss in ("l2", "l1", "cos") (residual_quantized.py:90) and has a dedicated l1 branch in _single_commitment_loss. The proto comment, the RQVAE.__init__ docstring (rqvae.py:52-53), and the in-code "flow step 3" comment (residual_quantized.py:297) all advertise only "l2"/"cos". l1 is a fully reachable, supported value — please document it everywhere or remove it from the assert.

Suggested change

	// Commitment loss type: "l2" or "cos".
	// Commitment loss type: "l2", "l1", or "cos".

[github-actions]

The proto declares optional string latent_weight = 11 [default = "1.0,0.5"] (sid_model.proto:54). In proto2 Python, accessing an unset field with a default returns the default literal — so cfg.latent_weight is always the truthy string "1.0,0.5" and the if cfg.latent_weight: guard is always True. The "let RQVAE / ResidualQuantized apply their signature default" branch is unreachable.

Not a behavioral bug (both paths yield (1.0, 0.5)), but the comment misrepresents control flow. If the intent really is "only override when user set it explicitly", drop the proto default and use cfg.HasField("latent_weight"). Otherwise simplify the comment and just always forward the value.

[github-actions]

self._is_inference is assigned False in __init__ (line 106) and never written anywhere else — there is no setter and RQVAE subclasses nn.Module directly, not BaseModule. The dispatch reduces to if not self.use_clip:. SidRqvae already calls forward_rqvae/forward_mixed directly and doesn't reach into RQVAE._is_inference, so it works "by accident".

Either drop the dead _is_inference flag and the misleading dispatch branch, or actually plumb BaseModule.set_is_inference through to the inner RQVAE so the documented behavior is real.

[github-actions]

Two concerns with the gather:

1. NCCL incompatibility. dist.gather_object requires a backend that supports object collectives. Production GPU training typically initialises a single NCCL group and gather_object raises RuntimeError: ProcessGroupNCCL does not support gather_object (PyTorch documents this). Consider initialising a side gloo group at module construction and passing it as group=, or replace with a tensor-based dist.gather after padding to a uniform per-rank shape.

2. Memory claim is off. The comment says peak rank-0 memory is O(world_size). gather_object pickles every rank's tensor into a Python bytes then deserialises into gathered, so rank 0 holds all world_size per-rank tensors plus the pickle bytestream simultaneously — closer to ~3× total corpus at peak. For large datasets this is a real OOM risk. Either chunk the gather (every rank ships fixed-size chunks in a loop) or write embeddings to per-rank shard files during predict and have rank 0 stream them.

[github-actions]

The output is forwarded straight through to faiss.Kmeans(D, K, **self.faiss_kmeans_kwargs) (residual_kmeans.py:244) with no whitelist or schema validation. A misspelled or unsupported kwarg (e.g. niters instead of niter, or gpu=True on a faiss-cpu build) raises inside faiss.Kmeans.__init__ only on rank 0, after gather_object has already pulled the entire corpus to rank 0. The other ranks then block on the subsequent broadcast/barrier until the NCCL watchdog tears the job down — no actionable error.

Two-step fix: (a) whitelist known FAISS keys (niter, nredo, seed, verbose, spherical, gpu, min/max_points_per_centroid, ...) and reject unknowns at construction time so the failure is symmetric across ranks; (b) construct the faiss.Kmeans object early (in __init__ or right before the gather) so config errors fire before the expensive cross-rank gather.

[github-actions]

B = Q.size(1) * dist.get_world_size() assumes uniform per-rank batch sizes. On the final training step (or any step with drop_last=False and uneven shards) this is wrong: each rank computes a different denominator at Q /= B (line 100) and a different scale at Q *= B (line 103). argmax is invariant to a positive per-row scalar so the final codes stay consistent — but the resulting Q values do not match across ranks, which biases any downstream code that relies on the assignment magnitudes.

Suggested fix: all_reduce the local Q.size(1) once at the top to get the true global B (one extra collective is negligible vs the ones already inside the loop). Also, dist.is_initialized() alone isn't sufficient — guard on dist.get_world_size() > 1 to avoid no-op collectives in single-process toy runs.

[github-actions]

Two issues that compound at scale:

1. Synchronous D2H every training step. embedding.detach().cpu() is a sync copy on a non-pinned host buffer — it blocks the CUDA stream until the copy completes. Use embedding.detach().to('cpu', non_blocking=True) against a pinned host buffer if you want overlap. Per-step impact is small but adds up.

2. Unbounded host-memory growth. The buffer grows linearly with steps × batch_size × D × 4 bytes per rank, and rank 0 later holds world_size × that during the gather (see the comment on the gather block). For a 100M-row dataset with D=512 you're looking at ~200 GB before the gather and ~world_size × that on rank 0 — a near-certain OOM at production scale.

Consider streaming per-step embeddings to a pre-allocated np.memmap of known max size (or per-rank chunked .npy files) and mmap-reading them at FAISS fit time. For very large datasets, reservoir-sampling at the buffer step is usually sufficient for K-Means quality (~256 × K samples is plenty).

[github-actions]

fallback = clip_mask.long().argmax() is computed from the local mask but safe_labels indexes into the global gathered logits (offset local_batch_size * self._rank + arange). When this rank has no clip rows, argmax returns 0 → safe_labels points at local column 0 of the global logits, which may itself be a recon column (all -inf). cross_entropy then yields NaN, which the nan_to_num at line 134-135 swallows before the mask multiply.

The current code is safe only because nan_to_num runs before * clip_mask.float() — 0 × NaN = NaN, so the ordering is load-bearing. Any later "optimization" that swaps the order, or that removes the seemingly-defensive nan_to_num, silently poisons the loss.

Suggested fix: compute fallback from the gathered mask (clip_mask_all.long().argmax()) so safe_labels always points at a real clip column globally; then the nan_to_num becomes unnecessary and the contract is explicit. Also consider reading dist.get_rank() lazily in forward instead of caching at __init__ (line 104) — modules can be constructed before init_process_group.

[github-actions]

Review summary — SID generation (RQ-VAE / RQ-KMeans)

Substantial, well-organized PR. The DDP on_train_end vote-then-bail logic, the persistent=False decision on _is_initialized (intent is correct even if the resume path is incomplete — see inline), the FAISS in-place residual updates, and the GatherLayer.backward slice-then-reduce trick are all good calls. Inline comments cover individual concerns; below are themes worth tracking at the PR level.

Cross-cutting issues to consider

• Checkpoint resume of SidRqkmeans is broken. Even with _is_initialized correctly non-persistent, there is no on_load hook that re-flips the flag after centroids are restored, so evaluate/predict against a saved checkpoint silently emits zero codes. See inline on kmeans.py:232.
• on_train_end lifecycle. Called only on the normal-exit path in _train_and_evaluate — if training raises (OOM, KeyboardInterrupt, dataloader error), the FAISS fit never runs but the docstring at sid_rqkmeans.py:17 says "called unconditionally." Either wrap in try/finally or update the docstring. Also worth deciding whether a failed/empty fit should block the terminal checkpoint write so users don't silently ship unusable checkpoints.
• Distributed code paths are untested. GatherLayer.backward, _sinkhorn's distributed reductions, and SidRqkmeans.on_train_end's DDP branch (gather_object/broadcast/empty-vote) all have zero test coverage — only single-process paths are exercised. A torch.multiprocessing.spawn smoke test on gloo with world_size=2 comparing MaskedCLIPLoss.forward → backward against a single-process reference would catch the most likely class of bug.
• Other untested branches: forward_mode="gumbel_softmax", rotation_trick=True, kmeans_init=True, distance_type="cosine", normalize_residuals=True, loss_type in {"l1","cosine"}, RQVAE.decode_codes/get_codes, and the ResidualKMeans.train_offline tensor-input path. Several test assertions are also "exists / requires_grad" rather than numerical (e.g. test_commitment_loss_l1_branch doesn't verify the l1 value actually differs from l2; test_rqvae_backward accepts any one non-zero grad anywhere). Adding setUp: torch.manual_seed(0) to both test classes would also remove a latent flakiness source.
• test_on_train_end_runs_faiss is silently skipped without faiss. If CI lacks FAISS the entire ResidualKMeans.train_offline (and KMeansLayer.load_centroids_) codepath becomes unverified while CI looks green. Either make FAISS a required test dep or add a fallback test using the pure-torch _residual_kmeans to validate the layered residual math.
• Sinkhorn perf in DDP. Each iteration of _sinkhorn issues one all_reduce(sum_of_rows), so per training step you pay n_layers × (1 + sinkhorn_iters) small NCCL collectives (18 with defaults). Mostly fixed launch overhead, but worth either batching across layers or documenting the cost.
• forward_mixed doubles encoder/decoder/quantizer cost for all rows in CLIP mode, including pure-recon rows whose fea2 pass is wasted (siamese semantics only require it for clip rows). Acceptable if CLIP batches are mostly clip pairs; worth documenting otherwise.

Doc/comment drift

• Proto comment, RQVAE.__init__ docstring, and ResidualQuantized.forward step-3 comment all advertise commitment_loss as "l2"|"cos", but the code accepts "l1" too (see inline on the proto).
• sid_rqvae.py:71-75 comment describes a fallback path that the proto default makes unreachable (see inline).
• gather_object block memory comment claims O(world_size) rank-0 memory — actual is closer to ~3× corpus due to pickle (see inline).

Nice work overall. The critical items are the checkpoint-resume gap, the gather_object NCCL/memory situation, and the missing distributed tests; everything else is cleanup.

[github-actions]

Correctness — Gumbel-Softmax does not propagate gradient into the encoder.

_compute_distances (and _find_nearest_embedding) are decorated with @torch.no_grad(), so distances returned here is fully detached. In the Gumbel branch:

weights = _gumbel_softmax_sample(-distances, temperature=temperature, hard=True)
emb = weights @ self.embedding.weight

weights is a function only of the detached distances, so gradient flows into self.embedding.weight but not back through x into the encoder. The STE branch happens to work because x + (q-x).detach() explicitly routes the gradient — Gumbel has no such shortcut.

Effect: choosing forward_mode="gumbel_softmax" silently trains only the codebook, not the encoder.

Fix: remove @torch.no_grad() from _compute_distances (the argmin/argmax step is naturally non-differentiable), or split out a _compute_distances_no_grad/_compute_distances_with_grad pair and use the differentiable one for the Gumbel logits.

[github-actions]

Correctness — averaging KMeans centroids across DDP ranks produces meaningless init.

_kmeans_plus_plus (kmeans.py:116–128) uses torch.randint/torch.multinomial with the per-rank default RNG, so each rank picks different seeds from its own local batch and runs Lloyd to different local minima. Cluster indices are unordered, so "centroid 0" on rank A and "centroid 0" on rank B correspond to different concepts. all_reduce(SUM)/world_size then averages permutation-misaligned centroids — the result is closer to noise than to a meaningful warm start.

The comment "every rank starts from the same codebook" is true but disguises that the codebook is junk.

Fix: run KMeans on rank 0 only, then dist.broadcast(c, src=0) — mirrors what SidRqkmeans.on_train_end already does for FAISS. Bonus: world_size× less compute.

[github-actions]

Perf — KMeans++ is O(K²·N), should be O(K·N).

Each loop iteration calls _squared_euclidean_distance(data, centroids[:i]), recomputing distances to all previously selected centroids. Standard KMeans++ tracks a running min_dists and only computes the distance to the newly selected centroid:

min_dists = _squared_euclidean_distance(data, centroids[:1]).squeeze(1)
for i in range(1, n_clusters):
    if min_dists.sum() == 0:
        centroids[i:] = data[torch.randint(0, N, (n_clusters - i,), device=data.device)]
        break
    next_idx = torch.multinomial(min_dists, num_samples=1)
    centroids[i] = data[next_idx]
    new_dists = _squared_euclidean_distance(data, centroids[i:i+1]).squeeze(1)
    min_dists = torch.minimum(min_dists, new_dists)

For K=256, N=batch_size this is ~256× more work than necessary on the first training step. Combined with _kmeans having no convergence early-stop (always 100 iters), the first-step startup stall under kmeans_init=True is substantial.

Also: torch.randint/torch.multinomial here use the global default generator with no seed — KMeans++ is non-deterministic across runs even with torch.manual_seed set elsewhere. Consider accepting/threading a torch.Generator.

[github-actions]

Perf — .item() on the hot path forces a GPU→host sync every forward.

is_initialized is read inside ResidualKMeans.forward (residual_kmeans.py:129) and all_initialized (residual_kmeans.py:101) on every eval batch, so each call serializes the CUDA stream. The flag is flipped exactly once (in load_centroids_ from on_train_end) — there's no reason to re-read the device tensor every step.

Fix: cache as a Python bool attribute, set in load_centroids_ and re-derived from the buffer in _load_from_state_dict. Keep _is_initialized as a buffer for checkpoint round-trip, but stop reading it on the forward path. For N_layers=3 this turns eval into 3 syncs per batch.

[github-actions]

Scalability — dist.gather_object of the full per-rank buffer will OOM rank 0 on realistic corpora.

Each rank concats its host buffer (line 282), then ships via gather_object (which pickle-serializes the tensor, ~2× transient bytes), then rank 0 concats again. Peak rank-0 memory is roughly 2 × world_size × per_rank_corpus × D × 4B. For a 10M-row × 64-D corpus across 8 ranks that's ~20 GB on rank 0 alone — taking down DDP training right at the end, after hours of compute.

Combined with the per-step .cpu() D2H (line 137, already TODO'd), the FAISS-only path is the largest operational risk in the PR.

Fix options:

1. Write per-rank buffers to a shared filesystem (torch.save(buffer, f"{model_dir}/sid_embed_rank{rank}.pt")), dist.barrier(), then have rank 0 stream them into a pre-allocated np.empty((total_N, D)) (size known via an all_gather of row counts first).
2. Reservoir-sample the per-rank buffer down to a configured cap (FAISS K-Means rarely benefits from > ~256 × K samples).

At minimum, expose a sample cap and add a guard that errors loudly when the gather would exceed available host memory.

[github-actions]

Comment contradicts behavior — if cfg.latent_weight: is always True.

The proto field has optional string latent_weight = 11 [default = "1.0,0.5"];, so cfg.latent_weight returns "1.0,0.5" (truthy) when unset and the user-provided string otherwise. The "let signature default apply" branch is unreachable; latent_weight is always forwarded.

If the intent is to detect "user explicitly set vs proto default", use cfg.HasField("latent_weight") — but that only works if you drop the proto-level default. Otherwise just always forward and delete the conditional + comment.

The same shape applies to commitment_loss (cfg defaults to "l2" per sid_model.proto:52) — the RQVAE.commitment_loss=None defaulting in rqvae.py:111 is dead from the proto-driven path.

[github-actions]

Stringly-typed list fields are inconsistent with the rest of tzrec.

hidden_dims, codebook, and latent_weight are optional string parsed at runtime via _sid_helpers.parse_int_list / parse_float_list. Everywhere else in tzrec these are repeated uint32 / repeated float (see optimizer.proto:226-228, tower.proto:158, most of feature.proto).

Cost of the string encoding:

• Defers all validation to model construction (int(x.strip()) raises ValueError on a typo instead of failing at proto load).
• Disables proto-level type checking, IDE completion, and text_format validation.
• Requires the ad-hoc _sid_helpers.py shim.
• The empty-string-truthy bug noted on sid_rqvae.py:67-71 would not exist with a repeated field (callers would always pass the list as-is).

Suggest repeated uint32 hidden_dims;, repeated uint32 codebook;, repeated float latent_weight;, then delete _sid_helpers.py.

Also: the comment "non-uniform codebooks such as 512,256,128 are supported" disagrees with ResidualKMeans.__init__ (residual_kmeans.py:82-85) which asserts uniformity for SidRqkmeans. Clarify the per-model rule.

[github-actions]

Perf — torch.unique(..., dim=0).shape[0] plus int / int triggers a GPU→host sync per eval batch.

torch.unique with dim=0 does a CUDA sort + dedupe; shape[0] then forces a host sync, and unique_sids / B (Python int / int) produces a Python float, syncing again. The same pattern is in sid_rqkmeans.py:229-230.

Fix:

# pack codes per-row into a single int64 — unique is much faster on 1-D
mults = torch.tensor([K ** i for i in range(codes.shape[1])], device=codes.device)
packed = (codes.long() * mults).sum(dim=1)
unique_sids = torch.unique(packed).numel()  # tensor-level until .update
self._metric_modules["unique_sid_ratio"].update(
    torch.tensor(unique_sids, device=codes.device, dtype=torch.float) / B
)

Or buffer packed codes across the eval epoch and run unique once at epoch end.

[github-actions]

Docstring drops 'l1'. ResidualQuantized asserts commitment_loss in ("l2", "l1", "cos") (residual_quantized.py:90), and sid_model.proto:51 also documents the three options.

Suggested change

	commitment_loss (str|None): commitment loss type ('l2'|'cos').
	commitment_loss (str|None): commitment loss type ('l2'|'l1'|'cos').

[github-actions]

Lifecycle ordering — on_train_end runs before the tail-save, but the in-loop checkpoint at last_ckpt_step was taken before the FAISS fit.

For SidRqkmeans, if the last in-loop save coincides with the final training step (last_ckpt_step == i_step), the tail-save block at lines 529-548 is skipped and the only persisted artifact has _is_initialized=False with zero centroids. Reloading that checkpoint silently yields a model that returns dummy-zero codes — KMeansLayer._load_from_state_dict's mid-fit guard does not fire because the flag is False (which reads as "uninitialized" rather than corrupted).

Either:

• Always force a final save after on_train_end() (independent of last_ckpt_step), or
• Have SidRqkmeans.on_train_end() overwrite the most recent checkpoint after the FAISS fit (only rank 0), or
• Refuse state_dict() serialization on KMeansLayer until _is_initialized=True.

[github-actions]

Review summary

Big, well-scoped PR. Models and module stack are cleanly separated, tests cover the happy paths for both SidRqvae and SidRqkmeans, and the on_train_end() hook is a reasonable way to land an offline-fit model in the existing train loop. Below is a consolidated punch list — inline comments have the full detail on each item.

Correctness (should block merge)

• Gumbel-Softmax does not propagate gradient into the encoder. _compute_distances is @torch.no_grad(), so the Gumbel branch only trains the codebook. STE works only by accident of its x + (q-x).detach() shortcut. See vector_quantize.py:250-254.
• DDP kmeans_init averages permutation-misaligned centroids. Each rank seeds KMeans++ independently and the all_reduce(SUM)/world_size blends unrelated clusters. Should be rank-0 fit + broadcast, like SidRqkmeans.on_train_end. See residual_quantized.py:183-188.
• on_train_end lifecycle vs. final-step checkpoint. If last_ckpt_step == i_step, the only persisted state is pre-FAISS-fit and the mid-fit corruption guard does not catch it. See main.py:514.

Performance / scalability

• KMeans++ is O(K²·N). Should incrementally update min_dists. Combined with _kmeans having no convergence early-stop (always 100 iters), this is a visible first-step stall. kmeans.py:119-128, 156-169.
• KMeansLayer.is_initialized .item() per forward. Per-batch GPU→host sync on the eval path; cache as a Python bool. kmeans.py:230-233.
• dist.gather_object of the full per-rank buffer. Rank-0 OOM risk at realistic corpus sizes; pickle adds a 2× transient. sid_rqkmeans.py:282-300. Combine with the already-TODO'd per-step .cpu() D2H (sid_rqkmeans.py:137) — the FAISS path is the largest operational risk.
• torch.unique(codes, dim=0).shape[0] host sync per eval batch. Pack codes to int64 first and keep arithmetic on-device. sid_rqvae.py:258, sid_rqkmeans.py:229-230.

Code quality / API consistency

• Protos use comma-separated string for hidden_dims / codebook / latent_weight while the rest of tzrec uses repeated fields. Drop _sid_helpers.py and use repeated uint32 / repeated float. sid_model.proto:38-54.
• latent_weight "default detection" comment contradicts behavior. Proto default is "1.0,0.5" (non-empty), so if cfg.latent_weight: is always True. Same shape applies to commitment_loss. sid_rqvae.py:67-71, rqvae.py:111-114.
• RQVAE.forward(*args, **kwargs) dispatch. Hides the signature; both call sites already invoke forward_rqvae / forward_mixed directly — consider removing the shim.

Test gaps worth filling

• No test exercises forward_mode="gumbel_softmax" — the buggy path above would never be caught by the current suite.
• No test exercises rotation_trick=True (25-line Householder math), distance_type="cosine", or loss_type ∈ {"l1","cosine"}.
• No multi-process test covers GatherLayer / _all_gather_with_grad / rank-offset labels in MaskedCLIPLoss — gloo @ world_size=2 on CPU would suffice.
• _kmeans_plus_plus non-determinism: no seed plumbed, so two identical-seed runs of the same model don't produce identical codebooks.

Documentation nits

• RQVAE docstring drops 'l1' from commitment_loss options.
• forward_mixed's "recon rows == fea1" annotation is not enforced and not used.
• _predict_mixed docstring omits the _is_inference fast-path.

Nice touches worth calling out: KMeansLayer._load_from_state_dict mid-fit poisoning guard, int32 (not bool) NCCL collectives in on_train_end, empty-rank deadlock avoidance via all_reduce(MAX) on the empty-flag, and the rank-0 fit + broadcast pattern in SidRqkmeans — the same pattern that would fix the DDP kmeans-init issue above.