NVIDIA-NeMo / NVIDIA-NeMo/RL

feat: extend checkpoint-engine refit to SGLang (and Megatron) generation backends

Open
#3,288 3 comments 0 reactions 0 assignees View on GitHub
Feature
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Summary

Checkpoint-engine refit (added in #2608) currently supports only the **vLLM** generation backend. Requests with a SGLang or Megatron generation backend are rejected at setup:

```python
# nemo_rl/algorithms/grpo.py
if checkpoint_engine_config is not None and backend != "vllm":
raise NotImplementedError(
"checkpoint-engine refit is only supported for the vLLM generation "
f"backend, but policy.generation.backend={backend!r}. ..."
)
```

This is a feature request to **extend checkpoint-engine refit (and the NIXL backend) to the other generation backends** — at minimum SGLang, and Megatron generation if applicable.

## Scope / clarification

The restriction is on the **generation** side only. Checkpoint-engine already supports **Megatron, DTensor, and DTensor v2 policy (training) backends** on the send side — so this request is specifically about the *receive* side (the generation engine) being limited to vLLM.

## Motivation

Users on SGLang rollouts cannot use the RDMA/NIXL non-colocated refit path and must fall back to the existing transports. A generation-backend-agnostic checkpoint-engine receive path would let SGLang (and any future backend) benefit from the same streamed, bucketed weight transfer.

## Notes

- The design doc already lists this under current limitations: "SGLang checkpoint-engine refit is not implemented." (`docs/design-docs/checkpoint-engines.md`).
- The `CheckpointEngine` transport abstraction is backend-agnostic; the work is on the generation-worker receive/apply side (an SGLang analog of `VllmCheckpointEngineMixin` / the worker-extension wiring).
- Tracking issue to be referenced from the setup `NotImplementedError` so users hitting the guard can follow progress (mirrors how #3275 tracks the PPO/distillation gap for sparse-delta refit).

_Filed as a follow-up to the #2608 review._

Contributor guide

Open the contributing guide

Research direction

Start in nemo_rl/algorithms/grpo.py at the generation-backend guard, then read docs/design-docs/checkpoint-engines.md and the VllmCheckpointEngineMixin and worker-extension wiring. Trace how the generation worker receives and applies checkpoint-engine transfers, and compare the existing vLLM path with the SGLang path. Done means SGLang refit works through the checkpoint-engine/NIXL receive path and the setup restriction and tracking reference are updated.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, distributed-systems
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.