Basic Tutorial: Adding a New LLM Inference/Serving Backend #21
Description
- Prerequisite: Make sure the LLM Inference framework can be launched following the SPMD style. For example, the LLM inference script can be launched by
torchrun --standalone --nproc=8 offline_inference.py - A Rollout class: Build a
xxx_rollout.pyscript similar tovllm_rollout.py. In this file, define axxxRolloutclass that inherits fromBaseRollout.- This class should have a
generate_sequenceAPI that accepts a batch ofinput_ids,response_masks, andposition_idsfrom theDataProtoas input. Theself.inference_engine(e.g.,LLMEnginein vLLM) is responsible for performing auto-regressive generation and outputting a batch of responses. These responses should then be concatenated withinput_ids, and theresponse_masksandposition_idsshould be reconstructed accordingly.
- This class should have a
- ShardingManager Classes for Weight Synchronization with Training Frameworks: Create files named
fsdp_xxx.pyandmegatron_xxx.py, similar tofsdp_vllm.pyandmegatron_vllm.py. These files should defineXXXShardingManagerclasses (i.e., HybridEngine) that handle weight sharding between the training and inference frameworks.- In
megatron_vllm.py, we define anAllGatherPPModelclass to collect weights across the pipeline parallelism dimension. The parameters stored in thememory_buffersofAllGatherPPModelwill be used to synchronize the weights with the models in the vLLM rollout.
- In
- Weight loading issues: It may be necessary to provide specific weight loaders for transferring weights between different LLM Inference and Training backends for each model. This is similar to the
dtensor_weight_loader.pyandmegatron_weight_loader.pyfiles in vLLM.