Skip to content

[Refactor] Reuse single-media pack tensors without torch.cat copies - #2239

Open
ShilohYu wants to merge 1 commit into
InternLM:mainfrom
ShilohYu:feat/glm53-media-cat
Open

ShilohYu wants to merge 1 commit into
InternLM:mainfrom
ShilohYu:feat/glm53-media-cat

Conversation

@ShilohYu

Copy link
Copy Markdown
Collaborator

Summary

Pure refactor (~40 lines of production code, plus 9 new unit tests). Three sites previously ran a full torch.cat copy even for single-element lists; they now reuse the same storage directly:

  • glm53_vl_tokenize_fn — single-image / single-video samples reuse the processor output
  • glm53_vl_sft_collator — single-media packs reuse the cast tensor
  • Glm53ForConditionalGeneration._prepare_llm_inputs — single-modality packs reuse the SequenceContext block

Reuse does not change any value, dtype, or shape. Vision payloads are read-only end-to-end (SP split, H2D, patch_embed), and the "single-element reuse / multi-element still cat" behavior is pinned by 9 new unit tests.

On top of #2207 (early bf16), this PR removes the second copy of every single-media pack.

Measured results

  • Microbenchmark (real training path): collate 9.24 → 4.62 ms (−50%), transient RSS delta 92.9 → 34.7 MB
  • End-to-end (single node 8×GPU, 512K / SP8 / GBS1 / 25B, 10 steps): per-rank peak RSS 110.6 → 86.3 GiB (−22.0%)
  • Numerical equivalence: step-1 loss bitwise identical across three runs; 10-step loss deviation below the same-code rerun noise floor; GPU memory bitwise unchanged (changes are CPU-side only)

@ShilohYu
ShilohYu force-pushed the feat/glm53-media-cat branch from cf6a2c9 to 76e9f66 Compare October 30, 2026 08:08

This branch has not been deployed

No deployments
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