(
build_config: BuildConfig,
rank: int = 0,
ckpt_dir: str = None,
model_config: Union[str, PretrainedConfig] = None,
model_cls=None,
dry_run:
bool = False, # return the modified BuildConfig without actually building the engine
**kwargs
)
| 297 | |
| 298 | |
| 299 | def build_model( |
| 300 | build_config: BuildConfig, |
| 301 | rank: int = 0, |
| 302 | ckpt_dir: str = None, |
| 303 | model_config: Union[str, PretrainedConfig] = None, |
| 304 | model_cls=None, |
| 305 | dry_run: |
| 306 | bool = False, # return the modified BuildConfig without actually building the engine |
| 307 | **kwargs |
| 308 | ) -> Union[Engine, BuildConfig]: |
| 309 | model_config = copy.deepcopy(model_config) |
| 310 | |
| 311 | logits_dtype = kwargs.get('logits_dtype') |
| 312 | if logits_dtype is not None: |
| 313 | model_config.logits_dtype = logits_dtype |
| 314 | |
| 315 | architecture = model_config.architecture |
| 316 | assert not build_config.plugin_config.streamingllm, \ |
| 317 | "StreamingLLM is no longer supported because attention sink cannot work with the non-cyclic kv cache kernel & runtime changes." |
| 318 | assert not build_config.plugin_config.pp_reduce_scatter or architecture == "MixtralForCausalLM", \ |
| 319 | "PP reduce scatter is only supported in the mixtral model." |
| 320 | |
| 321 | assert rank < model_config.mapping.world_size |
| 322 | |
| 323 | rank_config = copy.deepcopy(model_config) |
| 324 | rank_config.set_rank(rank) |
| 325 | |
| 326 | if model_cls is None: |
| 327 | assert architecture in MODEL_MAP, \ |
| 328 | f"Unsupported model architecture: {architecture}" |
| 329 | model_cls = MODEL_MAP[architecture] |
| 330 | if ckpt_dir is None: |
| 331 | model = model_cls(rank_config) |
| 332 | else: |
| 333 | model = model_cls.from_checkpoint(ckpt_dir, config=rank_config) |
| 334 | is_checkpoint_pruned = getattr(rank_config, 'is_pruned', False) |
| 335 | |
| 336 | if build_config.plugin_config.lora_plugin is not None: |
| 337 | lora_config = LoraConfig(lora_dir=kwargs['lora_dir'] or [], |
| 338 | lora_ckpt_source=kwargs['lora_ckpt_source'], |
| 339 | max_lora_rank=kwargs['max_lora_rank']) |
| 340 | if kwargs['lora_target_modules'] is not None: |
| 341 | # command line options is preferred over the modules in the lora dir |
| 342 | lora_config.lora_target_modules = kwargs['lora_target_modules'] |
| 343 | build_config.lora_config = lora_config |
| 344 | |
| 345 | if is_checkpoint_pruned or kwargs.pop('strip_plan', False): |
| 346 | build_config.use_strip_plan = True |
| 347 | build_config.use_refit = kwargs.get('refit', False) |
| 348 | |
| 349 | return build(model, build_config) |
| 350 | |
| 351 | |
| 352 | def build_and_save(rank, gpu_id, ckpt_dir, build_config, output_dir, log_level, |
no test coverage detected