feat: extend checkpoint-engine refit to SGLang (and Megatron) generation backends
- 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
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