执行摘要

修复了SLIME中权重更新过程中的CUDA IPC缓存泄漏，通过添加内存释放和CUDA IPC收集调用，避免GPU内存累积，提升系统长期运行的稳定性。

功能与动机

本PR旨在解决一个特定场景下的GPU内存泄漏问题。根据PR body，根因在于使用ForkingPickler时，GPU张量通过storage._share_cuda_()创建了CUDA IPC缓存条目，这些条目在消费者未关闭句柄前不会释放，导致内存累积。修复动机是确保内存能被及时回收，防止资源浪费。

实现拆解

修改集中在update_weight_from_tensor.py文件的update_weights函数中：

• 在循环中添加del long_lived_tensors, hf_named_tensors来明确释放GPU张量。
• 调用torch.cuda.ipc_collect()两次：一次在每次chunk处理后，释放已完成chunk的IPC缓存条目；一次在所有chunk完成后（barrier后），释放最后一个chunk的IPC条目。
  代码块示例：
  

  for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights):
      refs, long_lived_tensors = self._send_hf_params(hf_named_tensors)
      ray.get(refs)
      del long_lived_tensors, hf_named_tensors
      torch.cuda.ipc_collect()
      dist.barrier(group=get_gloo_group())
      torch.cuda.ipc_collect()
  

评论区精华

本PR未经过review讨论，直接由作者提交并合并，表明变更可能较为紧急或已在内部验证，无需额外争议。

风险与影响

风险： 引入的torch.cuda.ipc_collect()调用可能带来轻微性能开销，尤其是在高频率权重更新场景中。此外，依赖PyTorch CUDA IPC机制，需确保与现有环境的兼容性。
影响： 对用户透明地减少内存泄漏，改善系统稳定性；对开发团队，提供了一个针对CUDA内存管理的修复案例，有助于规避类似问题。

关联脉络

与近期PR 1765（涉及megatron_utils update_weight模块的优化）相关，表明团队正持续改进权重更新相关的性能和稳定性。建议结合阅读以理解该模块的整体演进方向。