(self, tags)
| 2919 | print_gpu_memory_use(f"After offloading memory [{tags}]", self.local_rank, self.device_id) |
| 2920 | |
| 2921 | def wakeup(self, tags): |
| 2922 | |
| 2923 | if tags == "weight" and self.use_cudagraph and self.is_kvcache_sleeping: |
| 2924 | raise RuntimeError( |
| 2925 | "Waking up [weight] alone is not supported when CUDA Graph is enabled, " |
| 2926 | "as recapturing the graph requires the KV cache to be rebuilt first. " |
| 2927 | "Please wake up [kv_cache] first." |
| 2928 | ) |
| 2929 | |
| 2930 | logger.info(f">>> start reloading memory, tags: {tags}") |
| 2931 | start_time = time.perf_counter() |
| 2932 | |
| 2933 | # Reset share_inputs to restore tensor shapes and values |
| 2934 | if self.speculative_method in ["mtp"]: |
| 2935 | self.proposer.model_inputs.reset_model_inputs() |
| 2936 | self.share_inputs.reset_share_inputs() |
| 2937 | |
| 2938 | # Reinitialize KV cache |
| 2939 | if "kv_cache" in tags.split(","): |
| 2940 | if not self.is_kvcache_sleeping: |
| 2941 | logger.info("GPU model runner's kv cache is not sleeping, no need to wakeup!") |
| 2942 | return |
| 2943 | if self.speculative_method in ["mtp"]: |
| 2944 | self.proposer.initialize_kv_cache(main_model_num_blocks=self.num_gpu_blocks) |
| 2945 | self.initialize_kv_cache() |
| 2946 | self.is_kvcache_sleeping = False |
| 2947 | |
| 2948 | # Reload weights, deepep_buffer, cudagraph, etc. |
| 2949 | if "weight" in tags.split(","): |
| 2950 | if not self.is_weight_sleeping: |
| 2951 | logger.info("GPU model runner's weight is not sleeping, no need to wakeup!") |
| 2952 | return |
| 2953 | if self.fd_config.parallel_config.shutdown_comm_group_if_worker_idle: |
| 2954 | self.dynamic_weight_manager.restart_communication_group() |
| 2955 | if self.fd_config.parallel_config.enable_expert_parallel: |
| 2956 | self.dynamic_weight_manager.recreate_deepep_buffer() |
| 2957 | self.dynamic_weight_manager.reload_model_weights() |
| 2958 | if self.use_cudagraph: |
| 2959 | self.capture_model() |
| 2960 | self.is_weight_sleeping = False |
| 2961 | |
| 2962 | logger.info(f"<<< finish reloading memory! time cost: {time.perf_counter()-start_time:.3f}s") |
| 2963 | print_gpu_memory_use(f"After reloading memory [{tags}]", self.local_rank, self.device_id) |
| 2964 | |
| 2965 | def padding_cudagraph_inputs(self) -> None: |
| 2966 | """ |
nothing calls this directly
no test coverage detected