Load weights_iterator.
(self, weights_iterator)
| 86 | return model_loader.load_model(fd_config=self.fd_config) |
| 87 | |
| 88 | def load_weights(self, weights_iterator): |
| 89 | """Load weights_iterator.""" |
| 90 | |
| 91 | context = paddle.LazyGuard() |
| 92 | architectures = f"{self.fd_config.model_config.architectures[0]}RL" |
| 93 | if self.fd_config.quant_config is not None: |
| 94 | quantization_context = multi_switch_config_context( |
| 95 | (self.fd_config.quant_config, "is_checkpoint_bf16", True), |
| 96 | (self.fd_config.load_config, "dynamic_load_weight", False), |
| 97 | ) |
| 98 | else: |
| 99 | # bf16 |
| 100 | quantization_context = multi_switch_config_context( |
| 101 | (self.fd_config.load_config, "dynamic_load_weight", False) |
| 102 | ) |
| 103 | with quantization_context: |
| 104 | with context: |
| 105 | model_cls = ModelRegistry.get_class(architectures) |
| 106 | self.tmp_model = model_cls(self.fd_config) |
| 107 | self.tmp_model.eval() |
| 108 | self.tmp_model.load_weights(weights_iterator) |
| 109 | if self.fd_config.speculative_config.model_type != "mtp": |
| 110 | process_final_after_loading(self.tmp_model, self.fd_config) |
| 111 | self.rollout_model = self.tmp_model |
| 112 | self.tmp_model = None |
| 113 | |
| 114 | def get_name_mappings_to_training(self, trainer_degree=None) -> Dict[str, str]: |
| 115 | """Get parameter name mappings between rollout and training models.""" |