(self, tags)
| 2886 | return self.dynamic_weight_manager.update_weights_by_rdma(version, verify_checksum) |
| 2887 | |
| 2888 | def sleep(self, tags): |
| 2889 | |
| 2890 | logger.info(f">>> start offloading memory, tags: {tags}") |
| 2891 | start_time = time.perf_counter() |
| 2892 | |
| 2893 | # Clear weights, deepep_buffer, cudagraph, etc. |
| 2894 | if "weight" in tags.split(","): |
| 2895 | if self.is_weight_sleeping: |
| 2896 | logger.info("GPU model runner's weight is already sleeping, no need to sleep again!") |
| 2897 | return |
| 2898 | if self.use_cudagraph: |
| 2899 | self.model.clear_grpah_opt_backend() |
| 2900 | if self.fd_config.parallel_config.enable_expert_parallel: |
| 2901 | self.dynamic_weight_manager.clear_deepep_buffer() |
| 2902 | self.dynamic_weight_manager.clear_model_weight() |
| 2903 | if self.fd_config.parallel_config.shutdown_comm_group_if_worker_idle: |
| 2904 | self.dynamic_weight_manager.clear_communication_group() |
| 2905 | self.is_weight_sleeping = True |
| 2906 | |
| 2907 | # Clear KV cache |
| 2908 | if "kv_cache" in tags.split(","): |
| 2909 | if self.is_kvcache_sleeping: |
| 2910 | logger.info("GPU model runner's kv cache is already sleeping, no need to sleep again!") |
| 2911 | return |
| 2912 | if self.speculative_method in ["mtp"]: |
| 2913 | self.proposer.clear_mtp_cache() |
| 2914 | self.clear_cache() |
| 2915 | self.is_kvcache_sleeping = True |
| 2916 | |
| 2917 | paddle.device.cuda.empty_cache() |
| 2918 | logger.info(f"<<< finish offloading memory! time cost: {time.perf_counter()-start_time:.3f}s") |
| 2919 | print_gpu_memory_use(f"After offloading memory [{tags}]", self.local_rank, self.device_id) |
| 2920 | |
| 2921 | def wakeup(self, tags): |
| 2922 |
no test coverage detected