()
| 467 | |
| 468 | |
| 469 | def main(): |
| 470 | parser = get_argument_parser() |
| 471 | args = parser.parse_args() |
| 472 | args_sanity_checks(args) |
| 473 | if args.dataset is None: |
| 474 | exp_name = f"{get_exp_name(args.config)}_{args.precision}_dummy_cal_{args.num_chunks}_chunks" |
| 475 | else: |
| 476 | exp_name = ( |
| 477 | f"{get_exp_name(args.config)}_{args.precision}_{args.num_chunks}_chunks" |
| 478 | ) |
| 479 | if args.platform == "DX4": |
| 480 | platform_b = b"mt6991" |
| 481 | elif args.platform == "DX3": |
| 482 | platform_b = b"mt6989" |
| 483 | else: |
| 484 | raise ValueError( |
| 485 | f"Platform should be either DX3 or DX4, but got {args.platform}" |
| 486 | ) |
| 487 | print_args(args, exp_name) |
| 488 | |
| 489 | config, weight_dir, tokenizer_class, chunk_class = resolve_model_classes( |
| 490 | args.config |
| 491 | ) |
| 492 | tokenizer = tokenizer_class.from_pretrained(weight_dir) |
| 493 | if args.preformatter is not None: |
| 494 | preformatter = Preformatter(args.preformatter) |
| 495 | |
| 496 | head_dim = int(config.head_dim) |
| 497 | |
| 498 | # Evenly distribute the layers across chunks. |
| 499 | num_blocks_per_chunk = [ |
| 500 | (config.num_hidden_layers // args.num_chunks) |
| 501 | + (i < (config.num_hidden_layers % args.num_chunks)) |
| 502 | for i in range(args.num_chunks) |
| 503 | ] |
| 504 | check_all_chunks_same_num_layer(num_blocks_per_chunk) # noqa: F405 |
| 505 | |
| 506 | output_folder = os.path.join("pte", exp_name) |
| 507 | |
| 508 | # Load all collected checkpoint files into one giant state_dict |
| 509 | state_dict = load_checkpoints(weight_dir) |
| 510 | |
| 511 | dump_embedding_lut_for_cmdline(weight_dir, state_dict, config) |
| 512 | |
| 513 | export_shapes, max_num_token, max_cache_size = get_export_shapes(args.shapes) |
| 514 | print(f"export shapes: {export_shapes}") |
| 515 | print(f"Max Num Token: {max_num_token}") |
| 516 | print(f"Max Cache Size: {max_cache_size}") |
| 517 | |
| 518 | if args.dataset is not None: |
| 519 | embedding_layer = get_embedding_layer(config, weight_dir, state_dict) |
| 520 | |
| 521 | # Instantiate model chunks |
| 522 | print("Instantiating submodels") |
| 523 | models = [] |
| 524 | for chunk_idx, num_blocks in enumerate(num_blocks_per_chunk): |
| 525 | chunk = chunk_class( |
| 526 | config, |
no test coverage detected