Skip to content

[docker] update torch memory saver - #1916

Merged
zhuzilin merged 1 commit into
mainfrom
zilin/tms
May 18, 2026
Merged

[docker] update torch memory saver#1916
zhuzilin merged 1 commit into
mainfrom
zilin/tms

Conversation

@zhuzilin

Copy link
Copy Markdown
Contributor

No description provided.

@zhuzilin
zhuzilin merged commit 2091f7a into main May 18, 2026
2 checks passed
@zhuzilin
zhuzilin deleted the zilin/tms branch May 18, 2026 04:15
aoshen02 added a commit to vllm-project/vime that referenced this pull request Jun 1, 2026
…emory (slime #1916, #1924, #1932)

Co-merge of build-coupled upstream slime PRs:
- THUDM/slime#1916 ([docker] update torch memory saver): bump the
  torch_memory_saver pin d64a639 -> a193d9dd1b877d33c64a41cfb3db9f867df2d926.
- THUDM/slime#1924 (Reduce host memory with upgraded tms): add TMS_CUDA_MAJOR
  build plumbing + the disable_grad_buffers_cpu_backup host-memory feature.
- THUDM/slime#1932 (save host memory): place the
  `offload_train -> disable_grad_buffers_cpu_backup` assignment AFTER the
  `use_critic -> offload_train=True` block, so critic runs (which enable
  offload_train) also get grad-buffer cpu-backup disabled. Folded into this
  batch because it directly fixes #1924's placement (cannot apply to main
  without #1924's line present).

Why co-merged: the a193d9dd tms wheel's setup.py raises RuntimeError on CUDA
builds unless TMS_CUDA_MAJOR is set (setup.py:117-124). The pin bump (#1916)
therefore cannot build without #1924's plumbing.

Changes:
- docker/Dockerfile: ARG TMS_CUDA_MAJOR + tms pin a193d9dd, default CUDA major from torch.
- docker/patch/latest/megatron.patch: +3 hunks (DDP _make_no_backup_context,
  param_and_grad_buffer no-backup region, training/training.py dp_extra_kwargs).
- docker/version.txt: nightly-dev-20260519a.
- vime/ray/actor_group.py: .so loader tries _cu12-suffixed name then unsuffixed.
- vime/backends/megatron_utils/update_weight/common.py: get_cpu_backup(x, zero_copy=True).
- vime/utils/arguments.py: --record-memory-history; offload_train implies
  disable_grad_buffers_cpu_backup (placed after use_critic per #1932).

Refs: THUDM/slime#1916, THUDM/slime#1924, THUDM/slime#1932, #107

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
aoshen02 added a commit to vllm-project/vime that referenced this pull request Jun 1, 2026
…emory (slime #1916, #1924, #1932)

Co-merge of build-coupled upstream slime PRs:
- THUDM/slime#1916 ([docker] update torch memory saver): bump the
  torch_memory_saver pin d64a639 -> a193d9dd1b877d33c64a41cfb3db9f867df2d926.
- THUDM/slime#1924 (Reduce host memory with upgraded tms): add TMS_CUDA_MAJOR
  build plumbing + the disable_grad_buffers_cpu_backup host-memory feature.
- THUDM/slime#1932 (save host memory): place the
  `offload_train -> disable_grad_buffers_cpu_backup` assignment AFTER the
  `use_critic -> offload_train=True` block, so critic runs (which enable
  offload_train) also get grad-buffer cpu-backup disabled. Folded into this
  batch because it directly fixes #1924's placement (cannot apply to main
  without #1924's line present).

Why co-merged: the a193d9dd tms wheel's setup.py raises RuntimeError on CUDA
builds unless TMS_CUDA_MAJOR is set (setup.py:117-124). The pin bump (#1916)
therefore cannot build without #1924's plumbing.

Changes:
- docker/Dockerfile: ARG TMS_CUDA_MAJOR + tms pin a193d9dd, default CUDA major from torch.
- docker/patch/latest/megatron.patch: +3 hunks (DDP _make_no_backup_context,
  param_and_grad_buffer no-backup region, training/training.py dp_extra_kwargs).
- docker/version.txt: nightly-dev-20260519a.
- vime/ray/actor_group.py: .so loader tries _cu12-suffixed name then unsuffixed.
- vime/backends/megatron_utils/update_weight/common.py: get_cpu_backup(x, zero_copy=True).
- vime/utils/arguments.py: --record-memory-history; offload_train implies
  disable_grad_buffers_cpu_backup (placed after use_critic per #1932).

Refs: THUDM/slime#1916, THUDM/slime#1924, THUDM/slime#1932, #107

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
CalvinXKY pushed a commit to vllm-project/vime that referenced this pull request Jun 2, 2026
…emory (slime #1916, #1924, #1932) (#114)

Co-merge of build-coupled upstream slime PRs:
- THUDM/slime#1916 ([docker] update torch memory saver): bump the
  torch_memory_saver pin d64a639 -> a193d9dd1b877d33c64a41cfb3db9f867df2d926.
- THUDM/slime#1924 (Reduce host memory with upgraded tms): add TMS_CUDA_MAJOR
  build plumbing + the disable_grad_buffers_cpu_backup host-memory feature.
- THUDM/slime#1932 (save host memory): place the
  `offload_train -> disable_grad_buffers_cpu_backup` assignment AFTER the
  `use_critic -> offload_train=True` block, so critic runs (which enable
  offload_train) also get grad-buffer cpu-backup disabled. Folded into this
  batch because it directly fixes #1924's placement (cannot apply to main
  without #1924's line present).

Why co-merged: the a193d9dd tms wheel's setup.py raises RuntimeError on CUDA
builds unless TMS_CUDA_MAJOR is set (setup.py:117-124). The pin bump (#1916)
therefore cannot build without #1924's plumbing.

Changes:
- docker/Dockerfile: ARG TMS_CUDA_MAJOR + tms pin a193d9dd, default CUDA major from torch.
- docker/patch/latest/megatron.patch: +3 hunks (DDP _make_no_backup_context,
  param_and_grad_buffer no-backup region, training/training.py dp_extra_kwargs).
- docker/version.txt: nightly-dev-20260519a.
- vime/ray/actor_group.py: .so loader tries _cu12-suffixed name then unsuffixed.
- vime/backends/megatron_utils/update_weight/common.py: get_cpu_backup(x, zero_copy=True).
- vime/utils/arguments.py: --record-memory-history; offload_train implies
  disable_grad_buffers_cpu_backup (placed after use_critic per #1932).

Refs: THUDM/slime#1916, THUDM/slime#1924, THUDM/slime#1932, #107

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
momo609 pushed a commit to vllm-project/vime that referenced this pull request Jun 8, 2026
…emory (slime #1916, #1924, #1932) (#114)

Co-merge of build-coupled upstream slime PRs:
- THUDM/slime#1916 ([docker] update torch memory saver): bump the
  torch_memory_saver pin d64a639 -> a193d9dd1b877d33c64a41cfb3db9f867df2d926.
- THUDM/slime#1924 (Reduce host memory with upgraded tms): add TMS_CUDA_MAJOR
  build plumbing + the disable_grad_buffers_cpu_backup host-memory feature.
- THUDM/slime#1932 (save host memory): place the
  `offload_train -> disable_grad_buffers_cpu_backup` assignment AFTER the
  `use_critic -> offload_train=True` block, so critic runs (which enable
  offload_train) also get grad-buffer cpu-backup disabled. Folded into this
  batch because it directly fixes #1924's placement (cannot apply to main
  without #1924's line present).

Why co-merged: the a193d9dd tms wheel's setup.py raises RuntimeError on CUDA
builds unless TMS_CUDA_MAJOR is set (setup.py:117-124). The pin bump (#1916)
therefore cannot build without #1924's plumbing.

Changes:
- docker/Dockerfile: ARG TMS_CUDA_MAJOR + tms pin a193d9dd, default CUDA major from torch.
- docker/patch/latest/megatron.patch: +3 hunks (DDP _make_no_backup_context,
  param_and_grad_buffer no-backup region, training/training.py dp_extra_kwargs).
- docker/version.txt: nightly-dev-20260519a.
- vime/ray/actor_group.py: .so loader tries _cu12-suffixed name then unsuffixed.
- vime/backends/megatron_utils/update_weight/common.py: get_cpu_backup(x, zero_copy=True).
- vime/utils/arguments.py: --record-memory-history; offload_train implies
  disable_grad_buffers_cpu_backup (placed after use_critic per #1932).

Refs: THUDM/slime#1916, THUDM/slime#1924, THUDM/slime#1932, #107

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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