deepseek-ai/Janus

Fork 0

mirror of https://github.com/deepseek-ai/Janus.git synced 2026-11-05 00:57:18 +00:00

[PR #232] [CLOSED] fix: use actual codebook size for usage tracking buffer in VectorQuantizer #229

New issue

Closed

opened 2026-10-26 21:14:21 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:14:21 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/Janus/pull/232
Author: @Mr-Neutr0n
Created: 2/11/2026
Status: ❌ Closed

Base: main ← Head: fix/vq-codebook-buffer-size

📝 Commits (1)

e24472a fix: use actual codebook size for usage buffer and remove nn.Parameter wrapper

📊 Changes

1 file changed (+3 additions, -2 deletions)

View changed files

📝 janus/models/vq_model.py (+3 -2)

📄 Description

Summary

Fixes two bugs in VectorQuantizer.__init__ in janus/models/vq_model.py:

1. Hardcoded buffer size (65536) instead of actual codebook size (`self.n_e`)

The codebook_used buffer was hardcoded to size 65536, but the actual codebook size (n_e) is configurable (default 16384). This causes:

Wasted memory when n_e < 65536 — the buffer allocates 4x more entries than needed at the default size.
Potential index-out-of-bounds errors if n_e > 65536 — usage tracking would fail for codebook entries beyond index 65535.
Misleading usage statistics — the buffer dimensions don't match the actual codebook, making usage tracking unreliable.

2. `nn.Parameter` wrapped inside `register_buffer`

register_buffer is designed to store tensors as non-parameter state (not updated by optimizers, but included in state_dict). Wrapping the tensor in nn.Parameter inside register_buffer creates conflicting semantics — nn.Parameter signals "this should be optimized," while register_buffer signals the opposite. A plain torch.zeros(...) tensor is the correct argument.

Fix

# Before
self.register_buffer("codebook_used", nn.Parameter(torch.zeros(65536)))

# After
self.register_buffer("codebook_used", torch.zeros(self.n_e))

This ensures the buffer size matches the actual codebook size and uses the correct tensor type for a registered buffer.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/Janus/pull/232 **Author:** [@Mr-Neutr0n](https://github.com/Mr-Neutr0n) **Created:** 2/11/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `fix/vq-codebook-buffer-size` --- ### 📝 Commits (1) - [`e24472a`](https://github.com/deepseek-ai/Janus/commit/e24472ae4d58ecd97ccf23feaea2fddfaa18a2ea) fix: use actual codebook size for usage buffer and remove nn.Parameter wrapper ### 📊 Changes **1 file changed** (+3 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `janus/models/vq_model.py` (+3 -2) </details> ### 📄 Description ## Summary Fixes two bugs in `VectorQuantizer.__init__` in `janus/models/vq_model.py`: ### 1. Hardcoded buffer size (65536) instead of actual codebook size (`self.n_e`) The `codebook_used` buffer was hardcoded to size `65536`, but the actual codebook size (`n_e`) is configurable (default `16384`). This causes: - **Wasted memory** when `n_e < 65536` — the buffer allocates 4x more entries than needed at the default size. - **Potential index-out-of-bounds errors** if `n_e > 65536` — usage tracking would fail for codebook entries beyond index 65535. - **Misleading usage statistics** — the buffer dimensions don't match the actual codebook, making usage tracking unreliable. ### 2. `nn.Parameter` wrapped inside `register_buffer` `register_buffer` is designed to store tensors as non-parameter state (not updated by optimizers, but included in `state_dict`). Wrapping the tensor in `nn.Parameter` inside `register_buffer` creates conflicting semantics — `nn.Parameter` signals "this should be optimized," while `register_buffer` signals the opposite. A plain `torch.zeros(...)` tensor is the correct argument. ### Fix ```python # Before self.register_buffer("codebook_used", nn.Parameter(torch.zeros(65536))) # After self.register_buffer("codebook_used", torch.zeros(self.n_e)) ``` This ensures the buffer size matches the actual codebook size and uses the correct tensor type for a registered buffer. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 21:14:21 +00:00

closed this issue
added the
pull-request
label

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

deepseek-ai/Janus#229

No description provided.

Rows
Columns