Support on the fly quant for rollout

Summary:
Only supporting quantizing all linear layers with torchao config for now. see vllm PR for
how to generate the quantization file.
Also requires vllm changes: vllm-project/vllm#23014

Test Plan:
sh examples/ppo_trainer/run_deepseek7b_llm.sh

Reviewers:

Subscribers:

Tasks:

Tags:

Author: jerryzh168

examples/ppo_trainer/run_deepseek7b_llm.sh

@@ -22,6 +22,8 @@ python3 -m verl.trainer.main_ppo \
     actor_rollout_ref.rollout.tensor_model_parallel_size=4 \
     actor_rollout_ref.rollout.name=vllm \
     actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
+    actor_rollout_ref.rollout.quantization=torchao \
+    actor_rollout_ref.rollout.quantization_config_file=torchao_config.json \
     critic.optim.lr=1e-5 \
     critic.model.use_remove_padding=True \
     critic.model.path=deepseek-ai/deepseek-llm-7b-chat \

verl/trainer/config/rollout/rollout.yaml

@@ -268,3 +268,6 @@ profiler:
 
   # specific tool config
   tool_config: ${oc.select:actor_rollout_ref.actor.profiler.tool_config,null}
+
+quantization: null
+quantization_config_file: null

verl/workers/config/rollout.py

@@ -139,3 +139,6 @@ class RolloutConfig(BaseConfig):
     layered_summon: bool = False
 
     layer_name_map: dict = field(default_factory=dict)
+
+    quantization: Optional[str] = None
+    quantization_config_file: Optional[str] = None

verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py

@@ -194,6 +194,8 @@ def __init__(self, model_path: str, config: RolloutConfig, tokenizer, model_hf_c
             enable_prefix_caching=True,
             trust_remote_code=trust_remote_code,
             seed=config.get("seed", 0),
+            quantization=config.quantization,
+            hf_overrides={"quantization_config": {"quantization_config_file": config.quantization_config_file}},
             **compilation_config,
             **lora_kwargs,
             **engine_kwargs,

verl/workers/sharding_manager/fsdp_vllm.py

@@ -338,16 +338,33 @@ def replace_lora_wrapper(k):
 
                 updated_params = {replace_lora_wrapper(k): v for k, v in updated_params.items()}
 
+
         from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader
 
         patch_vllm_moe_model_weight_loader(model)
         device = get_device_id()  # used when fsdp2 set cpu_offload_policy
-        loaded_params = model.load_weights(
-            (
-                (name, param.to(device, non_blocking=True).full_tensor() if isinstance(param, DTensor) else param)
-                for name, param in updated_params.items()
-            )
-        )
+
+        # make all DTensor full tensor before quantization
+        updated_params = {
+            name: param.to(device, non_blocking=True).full_tensor() if isinstance(param, DTensor) else param
+            for name, param in updated_params.items()
+        }
+
+        quantization = self.rollout_config.quantization
+        quantization_config_file = self.rollout_config.quantization_config_file
+        quantized_updated_params = {}
+        from vllm.model_executor.layers.quantization import get_quantization_config
+        import json
+        if quantization is not None and quantization_config_file is not None:
+            quant_cls = get_quantization_config(quantization)
+            config = quant_cls.from_config_file(quantization_config_file)
+            for name, param in updated_params.items():
+                if name.endswith("proj.weight"):
+                    quantized_updated_params[name] = config.quantize_param(param)
+                else:
+                    quantized_updated_params[name] = param
+
+        loaded_params = model.load_weights(quantized_updated_params.items())
 
         self.base_sync_done = True
         logger.info(f"vLLM load weights, loaded_params: {len(loaded_params) if loaded_params else -1}")