[RL] refactor flash rl weight reload in sglang (#14870)
Co-authored-by: eternally-z <zzywzj@gmail.com>
This commit is contained in:
@@ -788,6 +788,19 @@ class QuantizedRLModelLoader(DefaultModelLoader):
|
||||
"""
|
||||
logger.info("[QuantizedRL] Initial load with FP8 quantization")
|
||||
|
||||
original_load_weights = model.load_weights
|
||||
|
||||
def load_weights_proxy(weights):
|
||||
if QuantizedRLModelLoader.is_reload_scenario(model):
|
||||
logger.info("[QuantizedRL] Using fast path reload in load_weights")
|
||||
QuantizedRLModelLoader.rebinding_and_load_weights(
|
||||
model, original_load_weights, weights
|
||||
)
|
||||
else:
|
||||
original_load_weights(weights)
|
||||
|
||||
model.load_weights = load_weights_proxy
|
||||
|
||||
model.load_weights(weights)
|
||||
original_weights = dict(model.named_parameters())
|
||||
|
||||
|
||||
Reference in New Issue
Block a user