Segfault in checkpoint save (cudaMemcpyAsync) — regression in v0.2.4, works in v0.2.3

Description

Reproducible segfault during checkpoint save when using SLIME v0.2.4 with colocated mode and offload enabled. The same training configuration works without issues on SLIME v0.2.3.

Environment

• SLIME version: v0.2.4 (segfaults) vs v0.2.3 (works)
• PyTorch: 2.6
• GPU: 8× A100 40GB (single node)
• Model: Qwen3-4B
• Training: GRPO with --colocate --offload-rollout --offload-train

Steps to reproduce

1. Run any GRPO training job with --colocate and offload enabled
2. Let training complete a few rollout+train iterations
3. When checkpoint save triggers (either via --save-interval or at end of training), segfault occurs

The training loop itself runs correctly — rollout, tool execution, reward scoring, and GRPO loss computation all succeed. The crash only happens during save_model().

Stack trace

Context: saving checkpoint at iteration 4 in torch_dist format
         Storing distributed optimizer sharded state of type dp_reshardable

!!!!!!! Segfault encountered !!!!!!!

  cuMemcpyHtoDAsync_v2
  cudaMemcpyAsync
  at::native::copy_kernel_cuda(at::TensorIterator&, bool)
  at::native::copy_impl(at::Tensor&, at::Tensor const&, bool)
  at::native::copy_(at::Tensor&, at::Tensor const&, bool)
  at::_ops::copy_::call(at::Tensor&, at::Tensor const&, bool)
  at::native::_to_copy(...)
  ...autograd::VariableType::_to_copy(...)
  at::_ops::_to_copy::call(...)
  at::native::to(...)     ← tensor.to(device) in Python
  at::Tensor::to(c10::TensorOptions, ...) const
  torch::autograd::dispatch_to(...)
  torch::autograd::THPVariable_to(...)     ← Python tensor.to() call

Python traceback:
  train.py → train(args)
  slime/train.py:88 → save(rollout_id)
  slime/train.py:54 → actor_model.save_model(...)
  slime/ray/actor_group.py:133 → ray.get([actor.save_model.remote(...)])
  → ray.exceptions.ActorDiedError

All workers crash simultaneously with the same segfault.

Key observation

The segfault is in cuMemcpyHtoDAsync_v2 (Host-to-Device async copy), triggered by tensor.to(device) inside Megatron's checkpoint save path. This suggests that after the offload/reload cycles during training, some tensor data pointers become invalid, and the checkpoint save then tries to copy from/to stale pointers.

Bisect result

• v0.2.3: checkpoint save works correctly, no segfault
• v0.2.4: 100% reproducible segfault on checkpoint save

We have tried --no-pin-cpu-grads --no-pin-cpu-params as a workaround — these did not help in our case.

Related issues

• NVIDIA/Megatron-LM#1861 — similar segfault during async checkpoint save
• THUDM/slime#1786 — cudaError during offload_train (different root cause but same area)

Reproduce the failure with the stated SLIME v0.2.4 configuration, then compare the checkpoint path between v0.2.3 and v0.2.4. Start at slime/train.py:54 and :88 and slime/ray/actor_group.py:133, tracing save_model() through the distributed optimizer save; done means checkpoint saving completes without the CUDA segfault under colocated offload training.