From 2cdefa7d154a6b181a36c60de3c066d389d6ae6e Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Tue, 7 Oct 2025 08:04:05 +0000 Subject: [PATCH 1/3] refactor for torch 2.8 and vllm 0.10.2 --- flash_rl/fp8loader.py | 81 +++++++++++++++++++++++++++++++--------- flash_rl/vllm_patch.py | 84 +++++++++++++++++++++++++++++------------- 2 files changed, 122 insertions(+), 43 deletions(-) diff --git a/flash_rl/fp8loader.py b/flash_rl/fp8loader.py index 2dcf939..68791df 100644 --- a/flash_rl/fp8loader.py +++ b/flash_rl/fp8loader.py @@ -1,28 +1,73 @@ from vllm.device_allocator.cumem import CuMemAllocator from contextlib import contextmanager import torch -from torch.cuda.memory import MemPoolContext +# from torch.cuda.memory import MemPoolContext +from torch.cuda.memory import MemPool # new MemoryPool class still exported from torch._C import ( -_cuda_beginAllocateToPool, -_cuda_endAllocateCurrentStreamToPool, + # _cuda_beginAllocateToPool, + # _cuda_endAllocateCurrentStreamToPool, + _cuda_beginAllocateCurrentThreadToPool, + _cuda_endAllocateToPool, ) +# @contextmanager +# def disable_mem_pool(disable=False): +# if disable \ +# and "weights" in CuMemAllocator.get_instance().allocator_and_pools \ +# and MemPoolContext.active_pool() == \ +# CuMemAllocator.get_instance().allocator_and_pools["weights"][0]: +# pool = MemPoolContext.active_pool() +# ctx = MemPoolContext(None) +# device_index = torch.cuda.current_device() +# _cuda_endAllocateCurrentStreamToPool(device_index, pool.id) +# need_restart = True +# else: +# need_restart = False +# try: +# yield +# finally: +# if disable and need_restart: +# _cuda_beginAllocateToPool(device_index, pool.id) +# del ctx + + +def _get_pool_by_name(name: str) -> MemPool | None: + """Best-effort fetch of a named pool (e.g., 'weights') from vLLM.""" + pools = CuMemAllocator.get_instance().allocator_and_pools + entry = pools.get(name) + if not entry: + return None + # vLLM stored (pool, allocator?) tuples/lists before; keep robust. + cand = entry[0] + return cand if isinstance(cand, MemPool) else None + + @contextmanager -def disable_mem_pool(disable=False): - if disable \ - and "weights" in CuMemAllocator.get_instance().allocator_and_pools \ - and MemPoolContext.active_pool() == \ - CuMemAllocator.get_instance().allocator_and_pools["weights"][0]: - pool = MemPoolContext.active_pool() - ctx = MemPoolContext(None) - device_index = torch.cuda.current_device() - _cuda_endAllocateCurrentStreamToPool(device_index, pool.id) - need_restart = True - else: - need_restart = False +def disable_mem_pool(disable: bool = False, pool_name: str = "weights"): + """ + Temporarily stop routing allocations from *this thread* to the given MemPool. + On exit, restore routing to the same pool. + + Notes/pitfalls (new API): + - This only affects the current thread (new behavior). + - We cannot check the 'active' pool anymore (MemPoolContext removed), + so we optimistically end routing and then restore it. If routing + wasn't active, end is a no-op and restore is harmless. + """ + device_index = torch.cuda.current_device() + pool: MemPool | None = None + did_end = False + + if disable: + pool = _get_pool_by_name(pool_name) + if pool is not None: + # End routing of current thread to this pool (new API). + _cuda_endAllocateToPool(device_index, pool.id) + did_end = True + try: yield finally: - if disable and need_restart: - _cuda_beginAllocateToPool(device_index, pool.id) - del ctx + # Restore routing to the pool if we had ended it. + if disable and did_end and pool is not None: + _cuda_beginAllocateCurrentThreadToPool(device_index, pool.id) \ No newline at end of file diff --git a/flash_rl/vllm_patch.py b/flash_rl/vllm_patch.py index a08b2f4..90b8ffd 100644 --- a/flash_rl/vllm_patch.py +++ b/flash_rl/vllm_patch.py @@ -27,6 +27,25 @@ def bond_method_to_cls(func, obj): else: return types.MethodType(func, obj) +def safe_replace_param_data(param, new_data, preserved_attrs=None): + """Safely replace parameter data while preserving vLLM-specific attributes""" + if preserved_attrs is None: + preserved_attrs = ['tp_rank', 'tp_size', 'output_dim', 'input_dim', '_weight_loader'] + + # Save attributes before replacement + saved_attrs = {} + for attr in preserved_attrs: + if hasattr(param, attr): + saved_attrs[attr] = getattr(param, attr) + + # Replace data + param.data = new_data + + # Restore attributes if they were lost + for attr, value in saved_attrs.items(): + if not hasattr(param, attr): + setattr(param, attr, value) + def check_updated(name, updated_params, quant_fn_name): if name in updated_params: return True @@ -36,13 +55,12 @@ def check_updated(name, updated_params, quant_fn_name): and name[:-6] in updated_params: return True - return False + return False recorded_loader_keys = [ 'weight_loader', 'load_qkv_weight', 'load_row_parallel_weight', - 'load_column_parallel_weight', 'load_merged_column_weight', 'output_dim', 'input_dim', @@ -109,12 +127,7 @@ def hacked_process_weights_after_loading( quant_method = getattr(module, "quant_method", None) if isinstance(quant_method, QuantizeMethodBase): - if isinstance(quant_method, Fp8LinearMethod): - # for fast processing, we will do manual processing later - assert not quant_method.use_marlin, 'marlin (w8a16) does not support fp8_fast processing' - continue - - if isinstance(quant_method, CompressedTensorsW8A8Int8): + if isinstance(quant_method, Fp8LinearMethod) or isinstance(quant_method, CompressedTensorsW8A8Int8): # for fast processing, we will do manual processing later continue @@ -135,20 +148,20 @@ def hacked_process_weights_after_loading( ) pscale = all_updated_params[name + '_scale'] tmp_data = pscale.data - pscale.data = hacked_data_dict[name + '_scale'] + safe_replace_param_data(pscale, hacked_data_dict[name + '_scale']) del tmp_data else: strided_data = torch.as_strided( p.data, hacked_data_dict[name].shape, hacked_data_dict[name].stride()) hacked_data_dict[name].copy_(strided_data) tmp_data = p.data - p.data = hacked_data_dict[name] + safe_replace_param_data(p, hacked_data_dict[name]) del tmp_data else: skipped_params.append(name) tmp_data = p.data - p.data = hacked_data_dict[name] + safe_replace_param_data(p, hacked_data_dict[name]) del tmp_data else: assert 'int8' in model.flashrl_quant_fn, 'fast loading only supports int8 and fp8' @@ -162,7 +175,7 @@ def hacked_process_weights_after_loading( skipped_params.append(name) tmp_data = p.data - p.data = hacked_data_dict[name] + safe_replace_param_data(p, hacked_data_dict[name]) del tmp_data logger.debug(f"flash_rl load_weights skipped params: {skipped_params}") @@ -182,7 +195,7 @@ def hacked_process_weights_after_loading( skipped_params.append(name) tmp_data = p.data - p.data = hacked_data_dict[name] + safe_replace_param_data(p, hacked_data_dict[name]) del tmp_data logger.debug(f"flash_rl load_weights skipped params: {skipped_params}") @@ -482,7 +495,7 @@ def hacked_logprob_forward( sampled = greedy_sampled else: # Sampling - sampled = self.topk_topp_sampler( + sampled, _ = self.topk_topp_sampler( logits, sampling_metadata.generators, None, @@ -568,15 +581,9 @@ def hacked_init_( **kwargs ) -> None: - if parse(vllm.__version__) >= parse('0.10.1'): - logger.warning( - f'detected vLLM version {vllm.__version__}' - 'for vLLM > 0.10.1, `FlashRL` logprob patch has been upstreamed and thus skipped' - ) - else: - # Patch the sampler class - sampler_patch_status = patch_vllm_logprob_compute() - logger.debug(f"Patching vllm Sampler... status: {sampler_patch_status}") + # Patch the sampler class + sampler_patch_status = patch_vllm_logprob_compute() + logger.debug(f"Patching vllm Sampler... status: {sampler_patch_status}") import vllm.envs as envs assert envs.VLLM_USE_V1, 'flash_rl only supports vllm v1 for now' @@ -663,6 +670,20 @@ def hacked_init_( (config_data.get('fn', 'int8') != 'bf16'): model = vllm_model_finder(self) + + # CRITICAL FIX: Ensure all parameters have vLLM-required attributes + # In newer PyTorch/vLLM, parameters might be regular Parameters instead of BasevLLMParameter + # We need to add the missing attributes that vLLM's weight loaders expect + from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size + tp_rank = get_tensor_model_parallel_rank() + tp_size = get_tensor_model_parallel_world_size() + for name, param in model.named_parameters(): + if not hasattr(param, 'tp_rank'): + param.tp_rank = tp_rank + if not hasattr(param, 'tp_size'): + param.tp_size = tp_size + logger.debug(f"Added tp_rank={tp_rank} and tp_size={tp_size} to all model parameters") + quant_fn = config_data.get('fn', 'int8') model.flashrl_quant_fn = quant_fn logger.debug(f"flash_rl quantization function: {quant_fn}") @@ -697,7 +718,9 @@ def hacked_load_weights( for name, (shape, stride, dtype, nbytes) in model.hacked_original_weights_rebuild_keys.items(): if name in existing_params: - existing_params[name].data = torch.empty(shape, dtype=dtype) + param = existing_params[name] + new_data = torch.empty(shape, dtype=dtype, device=param.data.device) + safe_replace_param_data(param, new_data) for k, loader_k in model.hacked_recorded_loader.items(): for n, loader in loader_k.items(): @@ -738,7 +761,7 @@ def hacked_load_weights( skipped_params.append(name) tmp_data = p.data - p.data = hacked_data_dict[name] + safe_replace_param_data(p, hacked_data_dict[name]) del tmp_data logger.debug(f"flash_rl load_weights skipped params: {skipped_params}") @@ -756,6 +779,17 @@ def hacked_load_weights( setattr(module, attr, getattr(module, f'hacked_{attr}')) delattr(module, f'hacked_{attr}') + # CRITICAL: Ensure tp_rank/tp_size are on ALL parameters after weight manipulation + # Parameters may have lost these attributes during .data replacement + from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size + tp_rank = get_tensor_model_parallel_rank() + tp_size = get_tensor_model_parallel_world_size() + for name, param in model.named_parameters(): + if not hasattr(param, 'tp_rank'): + param.tp_rank = tp_rank + if not hasattr(param, 'tp_size'): + param.tp_size = tp_size + end_time = time.time() logger.debug(f"flash_rl load_weights process_weights_after_loading took {end_time - start_time:.2f} seconds") return updated_params From e2868b37d20457ffcf1003c901c152f736c17a60 Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Wed, 8 Oct 2025 08:47:16 +0000 Subject: [PATCH 2/3] add fixes --- flash_rl/fp8loader.py | 24 ------- flash_rl/vllm_patch.py | 149 ++++++++++++++++++++++++----------------- 2 files changed, 89 insertions(+), 84 deletions(-) diff --git a/flash_rl/fp8loader.py b/flash_rl/fp8loader.py index 68791df..2db4a59 100644 --- a/flash_rl/fp8loader.py +++ b/flash_rl/fp8loader.py @@ -1,36 +1,12 @@ from vllm.device_allocator.cumem import CuMemAllocator from contextlib import contextmanager import torch -# from torch.cuda.memory import MemPoolContext from torch.cuda.memory import MemPool # new MemoryPool class still exported from torch._C import ( - # _cuda_beginAllocateToPool, - # _cuda_endAllocateCurrentStreamToPool, _cuda_beginAllocateCurrentThreadToPool, _cuda_endAllocateToPool, ) -# @contextmanager -# def disable_mem_pool(disable=False): -# if disable \ -# and "weights" in CuMemAllocator.get_instance().allocator_and_pools \ -# and MemPoolContext.active_pool() == \ -# CuMemAllocator.get_instance().allocator_and_pools["weights"][0]: -# pool = MemPoolContext.active_pool() -# ctx = MemPoolContext(None) -# device_index = torch.cuda.current_device() -# _cuda_endAllocateCurrentStreamToPool(device_index, pool.id) -# need_restart = True -# else: -# need_restart = False -# try: -# yield -# finally: -# if disable and need_restart: -# _cuda_beginAllocateToPool(device_index, pool.id) -# del ctx - - def _get_pool_by_name(name: str) -> MemPool | None: """Best-effort fetch of a named pool (e.g., 'weights') from vLLM.""" pools = CuMemAllocator.get_instance().allocator_and_pools diff --git a/flash_rl/vllm_patch.py b/flash_rl/vllm_patch.py index 90b8ffd..c7c06bf 100644 --- a/flash_rl/vllm_patch.py +++ b/flash_rl/vllm_patch.py @@ -27,25 +27,6 @@ def bond_method_to_cls(func, obj): else: return types.MethodType(func, obj) -def safe_replace_param_data(param, new_data, preserved_attrs=None): - """Safely replace parameter data while preserving vLLM-specific attributes""" - if preserved_attrs is None: - preserved_attrs = ['tp_rank', 'tp_size', 'output_dim', 'input_dim', '_weight_loader'] - - # Save attributes before replacement - saved_attrs = {} - for attr in preserved_attrs: - if hasattr(param, attr): - saved_attrs[attr] = getattr(param, attr) - - # Replace data - param.data = new_data - - # Restore attributes if they were lost - for attr, value in saved_attrs.items(): - if not hasattr(param, attr): - setattr(param, attr, value) - def check_updated(name, updated_params, quant_fn_name): if name in updated_params: return True @@ -55,12 +36,13 @@ def check_updated(name, updated_params, quant_fn_name): and name[:-6] in updated_params: return True - return False + return False recorded_loader_keys = [ 'weight_loader', 'load_qkv_weight', 'load_row_parallel_weight', + 'load_column_parallel_weight', 'load_merged_column_weight', 'output_dim', 'input_dim', @@ -127,7 +109,12 @@ def hacked_process_weights_after_loading( quant_method = getattr(module, "quant_method", None) if isinstance(quant_method, QuantizeMethodBase): - if isinstance(quant_method, Fp8LinearMethod) or isinstance(quant_method, CompressedTensorsW8A8Int8): + if isinstance(quant_method, Fp8LinearMethod): + # for fast processing, we will do manual processing later + assert not quant_method.use_marlin, 'marlin (w8a16) does not support fp8_fast processing' + continue + + if isinstance(quant_method, CompressedTensorsW8A8Int8): # for fast processing, we will do manual processing later continue @@ -148,20 +135,20 @@ def hacked_process_weights_after_loading( ) pscale = all_updated_params[name + '_scale'] tmp_data = pscale.data - safe_replace_param_data(pscale, hacked_data_dict[name + '_scale']) + pscale.data = hacked_data_dict[name + '_scale'] del tmp_data else: strided_data = torch.as_strided( p.data, hacked_data_dict[name].shape, hacked_data_dict[name].stride()) hacked_data_dict[name].copy_(strided_data) tmp_data = p.data - safe_replace_param_data(p, hacked_data_dict[name]) + p.data = hacked_data_dict[name] del tmp_data else: skipped_params.append(name) tmp_data = p.data - safe_replace_param_data(p, hacked_data_dict[name]) + p.data = hacked_data_dict[name] del tmp_data else: assert 'int8' in model.flashrl_quant_fn, 'fast loading only supports int8 and fp8' @@ -175,7 +162,7 @@ def hacked_process_weights_after_loading( skipped_params.append(name) tmp_data = p.data - safe_replace_param_data(p, hacked_data_dict[name]) + p.data = hacked_data_dict[name] del tmp_data logger.debug(f"flash_rl load_weights skipped params: {skipped_params}") @@ -195,7 +182,7 @@ def hacked_process_weights_after_loading( skipped_params.append(name) tmp_data = p.data - safe_replace_param_data(p, hacked_data_dict[name]) + p.data = hacked_data_dict[name] del tmp_data logger.debug(f"flash_rl load_weights skipped params: {skipped_params}") @@ -495,7 +482,7 @@ def hacked_logprob_forward( sampled = greedy_sampled else: # Sampling - sampled, _ = self.topk_topp_sampler( + sampled = self.topk_topp_sampler( logits, sampling_metadata.generators, None, @@ -581,9 +568,15 @@ def hacked_init_( **kwargs ) -> None: - # Patch the sampler class - sampler_patch_status = patch_vllm_logprob_compute() - logger.debug(f"Patching vllm Sampler... status: {sampler_patch_status}") + if parse(vllm.__version__) >= parse('0.10.1'): + logger.warning( + f'detected vLLM version {vllm.__version__}' + 'for vLLM > 0.10.1, `FlashRL` logprob patch has been upstreamed and thus skipped' + ) + else: + # Patch the sampler class + sampler_patch_status = patch_vllm_logprob_compute() + logger.debug(f"Patching vllm Sampler... status: {sampler_patch_status}") import vllm.envs as envs assert envs.VLLM_USE_V1, 'flash_rl only supports vllm v1 for now' @@ -670,20 +663,6 @@ def hacked_init_( (config_data.get('fn', 'int8') != 'bf16'): model = vllm_model_finder(self) - - # CRITICAL FIX: Ensure all parameters have vLLM-required attributes - # In newer PyTorch/vLLM, parameters might be regular Parameters instead of BasevLLMParameter - # We need to add the missing attributes that vLLM's weight loaders expect - from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size - tp_rank = get_tensor_model_parallel_rank() - tp_size = get_tensor_model_parallel_world_size() - for name, param in model.named_parameters(): - if not hasattr(param, 'tp_rank'): - param.tp_rank = tp_rank - if not hasattr(param, 'tp_size'): - param.tp_size = tp_size - logger.debug(f"Added tp_rank={tp_rank} and tp_size={tp_size} to all model parameters") - quant_fn = config_data.get('fn', 'int8') model.flashrl_quant_fn = quant_fn logger.debug(f"flash_rl quantization function: {quant_fn}") @@ -718,9 +697,7 @@ def hacked_load_weights( for name, (shape, stride, dtype, nbytes) in model.hacked_original_weights_rebuild_keys.items(): if name in existing_params: - param = existing_params[name] - new_data = torch.empty(shape, dtype=dtype, device=param.data.device) - safe_replace_param_data(param, new_data) + existing_params[name].data = torch.empty(shape, dtype=dtype) for k, loader_k in model.hacked_recorded_loader.items(): for n, loader in loader_k.items(): @@ -733,6 +710,26 @@ def hacked_load_weights( logger.debug(f"flash_rl load_weights preparation took {end_time - start_time:.2f} seconds") start_time = end_time + # CRITICAL FIX: Ensure all parameters have vLLM-required attributes before weight loading + # This is needed for VLLM version compatibility after parameter manipulation + try: + from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size + tp_rank = get_tensor_model_parallel_rank() + tp_size = get_tensor_model_parallel_world_size() + + params_fixed = 0 + for name, param in model.named_parameters(): + if not hasattr(param, 'tp_rank'): + param.tp_rank = tp_rank + params_fixed += 1 + if not hasattr(param, 'tp_size'): + param.tp_size = tp_size + + if params_fixed > 0: + logger.debug(f"FLASH_RL: Added tp_rank/tp_size attributes to {params_fixed} parameters for VLLM compatibility") + except Exception as e: + logger.warning(f"FLASH_RL: Could not set tp_rank/tp_size attributes: {e}") + updated_params = original_load_weights( flash_quantize_fn(weights, self.flash_rl_profile) ) @@ -761,12 +758,32 @@ def hacked_load_weights( skipped_params.append(name) tmp_data = p.data - safe_replace_param_data(p, hacked_data_dict[name]) + p.data = hacked_data_dict[name] del tmp_data logger.debug(f"flash_rl load_weights skipped params: {skipped_params}") del skipped_params + # CRITICAL FIX: Restore tp_rank/tp_size attributes after parameter data manipulation + # Parameter .data replacement can cause loss of VLLM-specific attributes + try: + from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size + tp_rank = get_tensor_model_parallel_rank() + tp_size = get_tensor_model_parallel_world_size() + + params_fixed = 0 + for name, param in model.named_parameters(): + if not hasattr(param, 'tp_rank'): + param.tp_rank = tp_rank + params_fixed += 1 + if not hasattr(param, 'tp_size'): + param.tp_size = tp_size + + if params_fixed > 0: + logger.debug(f"FLASH_RL: Restored tp_rank/tp_size attributes to {params_fixed} parameters after data manipulation") + except Exception as e: + logger.warning(f"FLASH_RL: Could not restore tp_rank/tp_size attributes: {e}") + del hacked_data_dict gc.collect() torch.cuda.empty_cache() @@ -779,17 +796,6 @@ def hacked_load_weights( setattr(module, attr, getattr(module, f'hacked_{attr}')) delattr(module, f'hacked_{attr}') - # CRITICAL: Ensure tp_rank/tp_size are on ALL parameters after weight manipulation - # Parameters may have lost these attributes during .data replacement - from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size - tp_rank = get_tensor_model_parallel_rank() - tp_size = get_tensor_model_parallel_world_size() - for name, param in model.named_parameters(): - if not hasattr(param, 'tp_rank'): - param.tp_rank = tp_rank - if not hasattr(param, 'tp_size'): - param.tp_size = tp_size - end_time = time.time() logger.debug(f"flash_rl load_weights process_weights_after_loading took {end_time - start_time:.2f} seconds") return updated_params @@ -838,6 +844,30 @@ def test_reload_at_init_( device = torch.cuda.current_device() print(f"FLASH_RL re-loading model {config_i} to device {device}") model = vllm_model_finder(self) + + # CRITICAL FIX: Ensure all parameters have vLLM-required attributes + # In newer PyTorch/vLLM versions, parameters might be regular Parameters instead of BasevLLMParameter + # We need to add the missing attributes that vLLM's weight loaders expect + try: + from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size + tp_rank = get_tensor_model_parallel_rank() + tp_size = get_tensor_model_parallel_world_size() + + # First ensure existing model parameters have the required attributes + params_fixed = 0 + for name, param in model.named_parameters(): + if not hasattr(param, 'tp_rank'): + param.tp_rank = tp_rank + params_fixed += 1 + if not hasattr(param, 'tp_size'): + param.tp_size = tp_size + + if params_fixed > 0: + print(f"FLASH_RL: Added tp_rank/tp_size attributes to {params_fixed} parameters for VLLM compatibility") + except Exception as e: + print(f"FLASH_RL: Warning - Could not set tp_rank/tp_size attributes: {e}") + # Continue anyway, might still work + model.load_weights( ((name, param.to(device)) for name, param in model_to_be_reloaded.named_parameters()) ) @@ -857,4 +887,3 @@ def test_reload_at_init_( except Exception as e: logger.error(f"Error patching vllm reload LLM: {e}") return False - From 0c9af17958ab31f489796c8238cc19efe530772a Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Wed, 8 Oct 2025 09:08:14 +0000 Subject: [PATCH 3/3] simplify fixes --- flash_rl/vllm_patch.py | 66 ++---------------------------------------- 1 file changed, 2 insertions(+), 64 deletions(-) diff --git a/flash_rl/vllm_patch.py b/flash_rl/vllm_patch.py index c7c06bf..17c1f93 100644 --- a/flash_rl/vllm_patch.py +++ b/flash_rl/vllm_patch.py @@ -47,6 +47,8 @@ def check_updated(name, updated_params, quant_fn_name): 'output_dim', 'input_dim', '_assert_and_load', + 'tp_rank', + 'tp_size', ] def hacked_process_weights_after_loading( @@ -710,26 +712,6 @@ def hacked_load_weights( logger.debug(f"flash_rl load_weights preparation took {end_time - start_time:.2f} seconds") start_time = end_time - # CRITICAL FIX: Ensure all parameters have vLLM-required attributes before weight loading - # This is needed for VLLM version compatibility after parameter manipulation - try: - from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size - tp_rank = get_tensor_model_parallel_rank() - tp_size = get_tensor_model_parallel_world_size() - - params_fixed = 0 - for name, param in model.named_parameters(): - if not hasattr(param, 'tp_rank'): - param.tp_rank = tp_rank - params_fixed += 1 - if not hasattr(param, 'tp_size'): - param.tp_size = tp_size - - if params_fixed > 0: - logger.debug(f"FLASH_RL: Added tp_rank/tp_size attributes to {params_fixed} parameters for VLLM compatibility") - except Exception as e: - logger.warning(f"FLASH_RL: Could not set tp_rank/tp_size attributes: {e}") - updated_params = original_load_weights( flash_quantize_fn(weights, self.flash_rl_profile) ) @@ -764,26 +746,6 @@ def hacked_load_weights( logger.debug(f"flash_rl load_weights skipped params: {skipped_params}") del skipped_params - # CRITICAL FIX: Restore tp_rank/tp_size attributes after parameter data manipulation - # Parameter .data replacement can cause loss of VLLM-specific attributes - try: - from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size - tp_rank = get_tensor_model_parallel_rank() - tp_size = get_tensor_model_parallel_world_size() - - params_fixed = 0 - for name, param in model.named_parameters(): - if not hasattr(param, 'tp_rank'): - param.tp_rank = tp_rank - params_fixed += 1 - if not hasattr(param, 'tp_size'): - param.tp_size = tp_size - - if params_fixed > 0: - logger.debug(f"FLASH_RL: Restored tp_rank/tp_size attributes to {params_fixed} parameters after data manipulation") - except Exception as e: - logger.warning(f"FLASH_RL: Could not restore tp_rank/tp_size attributes: {e}") - del hacked_data_dict gc.collect() torch.cuda.empty_cache() @@ -844,30 +806,6 @@ def test_reload_at_init_( device = torch.cuda.current_device() print(f"FLASH_RL re-loading model {config_i} to device {device}") model = vllm_model_finder(self) - - # CRITICAL FIX: Ensure all parameters have vLLM-required attributes - # In newer PyTorch/vLLM versions, parameters might be regular Parameters instead of BasevLLMParameter - # We need to add the missing attributes that vLLM's weight loaders expect - try: - from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size - tp_rank = get_tensor_model_parallel_rank() - tp_size = get_tensor_model_parallel_world_size() - - # First ensure existing model parameters have the required attributes - params_fixed = 0 - for name, param in model.named_parameters(): - if not hasattr(param, 'tp_rank'): - param.tp_rank = tp_rank - params_fixed += 1 - if not hasattr(param, 'tp_size'): - param.tp_size = tp_size - - if params_fixed > 0: - print(f"FLASH_RL: Added tp_rank/tp_size attributes to {params_fixed} parameters for VLLM compatibility") - except Exception as e: - print(f"FLASH_RL: Warning - Could not set tp_rank/tp_size attributes: {e}") - # Continue anyway, might still work - model.load_weights( ((name, param.to(device)) for name, param in model_to_be_reloaded.named_parameters()) )