Initializes model parallelism for the given model and device map. Args: raw_model_dir (str): The directory containing the pre-trained model files. device_map (Union[str, List[int]], optional): The list of GPU device indices for model parallelism, or "auto" t
(self, raw_model_dir: str, device_map: Union[str, List[int]] = "auto")
| 85 | self.moss_stopwords = torch.LongTensor([self.tokenizer.convert_tokens_to_ids("<eom>")]) |
| 86 | |
| 87 | def Init_Model_Parallelism(self, raw_model_dir: str, device_map: Union[str, List[int]] = "auto") -> MossForCausalLM: |
| 88 | """ |
| 89 | Initializes model parallelism for the given model and device map. |
| 90 | |
| 91 | Args: |
| 92 | raw_model_dir (str): The directory containing the pre-trained model files. |
| 93 | device_map (Union[str, List[int]], optional): The list of GPU device indices for model parallelism, or "auto" to use the default device map. Defaults to "auto". |
| 94 | |
| 95 | Returns: |
| 96 | MossForCausalLM: The model with model parallelism initialized. |
| 97 | |
| 98 | References: |
| 99 | https://github1s.com/huggingface/accelerate/blob/HEAD/src/accelerate/big_modeling.py#L407 |
| 100 | """ |
| 101 | # Print the number of CUDA devices available |
| 102 | print("Model Parallelism Devices: ", torch.cuda.device_count()) |
| 103 | if not os.path.exists(raw_model_dir): |
| 104 | raw_model_dir = snapshot_download(raw_model_dir) |
| 105 | |
| 106 | # Load model configuration from the raw_model_dir |
| 107 | config = MossConfig.from_pretrained(raw_model_dir) |
| 108 | |
| 109 | # Initialize an empty model with the loaded configuration and set the data type to float16 |
| 110 | with init_empty_weights(): |
| 111 | raw_model = MossForCausalLM._from_config(config, torch_dtype=torch.float16) |
| 112 | |
| 113 | # Tie the model's weights |
| 114 | raw_model.tie_weights() |
| 115 | |
| 116 | # Load the checkpoint and dispatch the model to the specified devices |
| 117 | model = load_checkpoint_and_dispatch( |
| 118 | raw_model, |
| 119 | raw_model_dir, |
| 120 | device_map="auto" if not device_map else device_map, |
| 121 | no_split_module_classes=["MossBlock"], |
| 122 | dtype=torch.float16 |
| 123 | ) |
| 124 | |
| 125 | return model |
| 126 | |
| 127 | def preprocess(self, raw_text: str) -> Tuple[torch.Tensor, torch.Tensor]: |
| 128 | """ |