()
| 412 | |
| 413 | |
| 414 | def main(): |
| 415 | parser = get_argument_parser() |
| 416 | args = parser.parse_args() |
| 417 | args_sanity_checks(args) |
| 418 | if args.dataset is None: |
| 419 | exp_name = f"{get_exp_name(args.config)}_{args.precision}_dummy_cal_{args.num_chunks}_chunks" |
| 420 | else: |
| 421 | exp_name = ( |
| 422 | f"{get_exp_name(args.config)}_{args.precision}_{args.num_chunks}_chunks" |
| 423 | ) |
| 424 | if args.platform == "DX4": |
| 425 | platform_b = b"mt6991" |
| 426 | elif args.platform == "DX3": |
| 427 | platform_b = b"mt6989" |
| 428 | else: |
| 429 | raise ValueError( |
| 430 | f"Platform should be either DX3 or DX4, but got {args.platform}" |
| 431 | ) |
| 432 | print_args(args, exp_name) |
| 433 | |
| 434 | config, weight_dir, tokenizer_class, chunk_class = resolve_model_classes( |
| 435 | args.config |
| 436 | ) |
| 437 | tokenizer = tokenizer_class.from_pretrained(weight_dir) |
| 438 | if args.preformatter is not None: |
| 439 | preformatter = Preformatter(args.preformatter) |
| 440 | |
| 441 | head_dim = int(config.head_dim) |
| 442 | |
| 443 | # Evenly distribute the layers across chunks. |
| 444 | num_blocks_per_chunk = [ |
| 445 | (config.num_hidden_layers // args.num_chunks) |
| 446 | + (i < (config.num_hidden_layers % args.num_chunks)) |
| 447 | for i in range(args.num_chunks) |
| 448 | ] |
| 449 | check_all_chunks_same_num_layer(num_blocks_per_chunk) # noqa: F405 |
| 450 | |
| 451 | output_folder = os.path.join("pte", exp_name) |
| 452 | |
| 453 | # Load all collected checkpoint files into one giant state_dict |
| 454 | state_dict = load_checkpoints(weight_dir) |
| 455 | |
| 456 | dump_embedding_lut_for_cmdline(weight_dir, state_dict, config) |
| 457 | |
| 458 | export_shapes, max_num_token, max_cache_size = get_export_shapes(args.shapes) |
| 459 | print(f"export shapes: {export_shapes}") |
| 460 | print(f"Max Num Token: {max_num_token}") |
| 461 | print(f"Max Cache Size: {max_cache_size}") |
| 462 | |
| 463 | if args.dataset is not None: |
| 464 | embedding_layer = get_embedding_layer(config, weight_dir, state_dict) |
| 465 | |
| 466 | # Instantiate model chunks |
| 467 | print("Instantiating submodels") |
| 468 | models = [] |
| 469 | for chunk_idx, num_blocks in enumerate(num_blocks_per_chunk): |
| 470 | chunk = chunk_class( |
| 471 | config, |
no test coverage detected