Skip to content

Fix boolean tensor memory accounting - #4236

Open
Excelius-Wang wants to merge 1 commit into
huggingface:mainfrom
Excelius-Wang:fix-bool-buffer-memory
Open

Fix boolean tensor memory accounting#4236
Excelius-Wang wants to merge 1 commit into
huggingface:mainfrom
Excelius-Wang:fix-bool-buffer-memory

Conversation

@Excelius-Wang

Copy link
Copy Markdown

dtype_byte_size(torch.bool) currently returns 1/8, but PyTorch boolean tensors use one byte per element. This undercounts boolean buffers in module-size estimates and can place tensors within a max_memory budget that is smaller than their actual storage.

Return 1 for torch.bool, reusing the existing sizing and placement paths. Keep the packed CustomDtype estimates unchanged.

A small randomly initialized GPT-2 using Transformers 4.57.6 has two 1 MiB boolean attention buffers. Before this change, Accelerate reports 511,496 bytes for a model with 2,346,504 bytes of parameters and buffers; after the change, the estimate matches. This demonstrates the storage-accounting issue, not a measured production OOM. GPT-2 in Transformers 5.15.1 no longer has these registered buffers.

Validation:

  • New/updated regression selection: 15 failures and 8 passes before the fix; all pass after.
  • Modeling utilities: 57 passed, 9 skipped (plus 4 passing subtests).
  • Big modeling: 11 passed, 26 skipped on CPU.
  • Covers empty and non-empty buffers, persistence, dtype overrides, nested totals, and automatic placement at budget boundaries.
  • Full make quality and git diff --check pass.

Scope: this corrects boolean tensor accounting. The existing module-splitting path can separately omit a module's own buffers when no_split_module_classes is not supplied; this patch does not address that behavior or guarantee a hard runtime memory bound.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant