[Bug] Colocate weight update fails with torch_memory_saver/offload because PyTorch CUDA IPC _share_cuda_ raises cudaErrorInvalidValue

Aperta
  • Qualcosa non va?
#2,188 2 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
45/100
Tipo di issue
Bug
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
distributed-systems, machine-learning

Direzione di ricerca

Inizia da slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py e traccia il percorso co-localizzato da actor_model.update_weights() fino a MultiprocessingSerializer.serialize. Riproduci il primo aggiornamento dei pesi con le impostazioni indicate di colocate, offload-train, NCCL e torch_memory_saver, quindi verifica che i motori co-localizzati completino l’aggiornamento senza l’errore CUDA IPC e che il comportamento di trasporto configurato sia supportato o documentato.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

bug
Bug Description

I am running a Slime GRPO job with Megatron backend, SGLang rollout engines, --colocate, --offload-train, and torch_memory_saver preload enabled. The job consistently fails during the first colocated weight update, before any training step completes.

Although --update-weight-transport nccl is configured, colocated engines still go through UpdateWeightFromTensor, which serializes flattened CUDA tensors with MultiprocessingSerializer / PyTorch CUDA IPC. In this environment that path fails in storage._share_cuda_() with CUDA error: invalid argument.

Steps to Reproduce

  1. Run a single-node 8xH100 Slime GRPO job with:

    • Megatron backend
    • SGLang rollout engines
    • --colocate
    • --offload-train
    • --update-weight-mode full
    • --update-weight-transport nccl
    • --rollout-num-gpus 8
    • --rollout-num-gpus-per-engine 1
    • --tensor-model-parallel-size 4
    • --megatron-to-hf-mode bridge

  2. Enable torch_memory_saver preload:

    • LD_PRELOAD=<torch_memory_saver_hook_mode_preload_cu12.abi3.so>
    • TMS_INIT_ENABLE=1
    • TMS_INIT_ENABLE_CPU_BACKUP=1

  3. Start training.

  4. Observe that SGLang engines start successfully and Megatron loads the checkpoint, but the first actor_model.update_weights() fails inside UpdateWeightFromTensor.

Expected Behavior

Slime colocate should either:

  • successfully update colocated SGLang engines with CUDA IPC under torch_memory_saver, or
  • provide a supported non-disk fallback, such as local NCCL weight update for colocated engines.

In particular, if --update-weight-transport nccl is configured, it would be helpful to either use NCCL for colocated engines too or document that colocated engines still use PyTorch CUDA IPC.

Actual Behavior

The first colocated weight update fails in PyTorch CUDA IPC serialization:

ray.exceptions.RayTaskError(AcceleratorError): ray::MegatronTrainRayActor.update_weights()
  File ".../slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py", line ...
    serialized = MultiprocessingSerializer.serialize(flattened_tensor_data, output_str=True)
  File ".../sglang/srt/utils/common.py", line 2335, in serialize
    ForkingPickler(buf).dump(obj)
  File ".../torch/multiprocessing/reductions.py", line 353, in reduce_tensor
    ) = storage._share_cuda_()
  File ".../torch/storage.py", line 1452, in _share_cuda_
    return self._untyped_storage._share_cuda_(*args, **kwargs)
torch.AcceleratorError: CUDA error: invalid argument

The failing tensor can be very small:

rank=6 dtype=dtype tms_interesting_region=False
cuda_context={'current_device': 6, 'local_rank': '6',
              'cuda_visible_devices': '0,1,2,3,4,5,6,7'}
tensor={'shape': (8192,), 'dtype': 'torch.uint8', 'device': 'cuda:6',
        'numel': 8192, 'nbytes': 8192, 'is_contiguous': True, ...}

I also ruled out current-device mismatch and simple storage reuse:

  • current_device == local_rank == tensor.device
  • cloned tensors with a different storage pointer still fail in _share_cuda_() in the full Slime actor after the actor enters the TMS/offload lifecycle
  • cloning and serializing from a fresh Python thread also fails in _share_cuda_()
Environment
  • slime version: 0.3.0, commit 4c058f65
  • Python version: 3.12.3
  • PyTorch version: 2.11.0+cu129
  • CUDA/ROCm version: CUDA 12.9
  • OS: Ubuntu 24.04.2 LTS
  • SGLang version: 0.5.13, commit 28b095c
  • Megatron-LM version: commit 1dcf0dafa
  • Ray version: 2.55.1
  • torch_memory_saver version: 0.0.9.post1
Logs
ray.exceptions.RayTaskError(AcceleratorError): ray::MegatronTrainRayActor.update_weights()
  File ".../slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py", line ...
    serialized = MultiprocessingSerializer.serialize(flattened_tensor_data, output_str=True)
  File ".../sglang/srt/utils/common.py", line 2335, in serialize
    ForkingPickler(buf).dump(obj)
  File ".../torch/multiprocessing/reductions.py", line 353, in reduce_tensor
    ) = storage._share_cuda_()
  File ".../torch/storage.py", line 1452, in _share_cuda_
    return self._untyped_storage._share_cuda_(*args, **kwargs)
torch.AcceleratorError: CUDA error: invalid argument

CUDA IPC serialize failed for colocate weight bucket;
rank=6 dtype=dtype tms_interesting_region=False
cuda_context={'current_device': 6, 'local_rank': '6',
              'cuda_visible_devices': '0,1,2,3,4,5,6,7'}
tensor={'shape': (8192,), 'dtype': 'torch.uint8', 'device': 'cuda:6',
        'numel': 8192, 'nbytes': 8192, 'is_contiguous': True, ...}

CUDA IPC serialize retry failed for colocate weight bucket;
cloned_tensor={'shape': (8192,), 'dtype': 'torch.uint8', 'device': 'cuda:0',
               'data_ptr': 43298350592, 'storage_ptr': 43298350592}
thread_traceback=...
  File ".../torch/multiprocessing/reductions.py", line 353, in reduce_tensor
    ) = storage._share_cuda_()
torch.AcceleratorError: CUDA error: invalid argument
Additional Context

I checked nearby issues/PRs:

  • THUDM/slime#1731 fixes CUDA IPC cache lifetime/leak behavior, but this failure happens at the first _share_cuda_() call.
  • THUDM/slime#1786 is another TMS/offload CUDA error, but the stack is different.
  • THUDM/slime#2186 is about CUDA 13 selecting the wrong torch_memory_saver preload .so; my environment is CUDA 12.9, the preload library loads successfully, and the failure happens later during the first colocated weight update.
  • fzyzcjy/torch_memory_saver#36 reports a similar lower-level issue where CUDA IPC tensor serialization fails when torch_memory_saver is enabled.
Pre-submission Checklist
  • I have read the CONTRIBUTING.md and understand the collaboration scope.
  • I have read the documentation and my issue is not addressed there.
  • I have searched for existing issues and this is not a duplicate.
  • I have provided a minimal, reproducible example.
Lingua principale
Python
Stelle
8.8k
Fork
1.3k
Merge medio
20h 44m
PR unite (30g)
25

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di THUDM/slime

Tutte le issue di THUDM/slime

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.