Warp model. Args: model: Model instance or model name. **kwargs: Additional keyword arguments.
(self, model, **kwargs)
| 943 | torch.distributed.destroy_process_group() |
| 944 | |
| 945 | def warp_model(self, model, **kwargs): |
| 946 | |
| 947 | """Warp model. |
| 948 | |
| 949 | Args: |
| 950 | model: Model instance or model name. |
| 951 | **kwargs: Additional keyword arguments. |
| 952 | """ |
| 953 | if self.use_deepspeed: |
| 954 | from deepspeed.runtime.zero.stage_1_and_2 import ( |
| 955 | estimate_zero2_model_states_mem_needs_all_live, |
| 956 | ) |
| 957 | from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live |
| 958 | from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict |
| 959 | |
| 960 | local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", 1)) |
| 961 | world_size = int(os.environ.get("WORLD_SIZE", 1)) |
| 962 | |
| 963 | # NOTE(xcsong): look in detail how the memory estimator API works: |
| 964 | # https://deepspeed.readthedocs.io/en/latest/memory.html#discussion |
| 965 | if int(os.environ.get("RANK", 0)) == 0: |
| 966 | logging.info("Estimating model states memory needs (zero2)...") |
| 967 | estimate_zero2_model_states_mem_needs_all_live( |
| 968 | model, |
| 969 | num_gpus_per_node=local_world_size, |
| 970 | num_nodes=world_size // local_world_size, |
| 971 | ) |
| 972 | logging.info("Estimating model states memory needs (zero3)...") |
| 973 | estimate_zero3_model_states_mem_needs_all_live( |
| 974 | model, |
| 975 | num_gpus_per_node=local_world_size, |
| 976 | num_nodes=world_size // local_world_size, |
| 977 | ) |
| 978 | device = None # Init device later |
| 979 | pass # Init DeepSpeed later |
| 980 | |
| 981 | elif self.use_ddp: |
| 982 | local_rank = int(os.environ.get("LOCAL_RANK", 0)) |
| 983 | model = model.cuda(local_rank) |
| 984 | model = DDP( |
| 985 | model, |
| 986 | device_ids=[local_rank], |
| 987 | find_unused_parameters=kwargs.get("train_conf", {}).get( |
| 988 | "find_unused_parameters", False |
| 989 | ), |
| 990 | ) |
| 991 | |
| 992 | else: |
| 993 | model = model.to(device=kwargs.get("device", "cuda")) |
| 994 | |
| 995 | return model |
| 996 | |
| 997 | def warp_optim_scheduler(self, model, **kwargs): |
| 998 | """Warp optim scheduler. |