(
output_folder,
exp_name,
model,
precision,
max_num_token,
max_cache_size,
chunk_idx,
export_shapes,
platform_b,
cal_dataset=None,
)
| 377 | |
| 378 | |
| 379 | def export_to_et_ir( |
| 380 | output_folder, |
| 381 | exp_name, |
| 382 | model, |
| 383 | precision, |
| 384 | max_num_token, |
| 385 | max_cache_size, |
| 386 | chunk_idx, |
| 387 | export_shapes, |
| 388 | platform_b, |
| 389 | cal_dataset=None, |
| 390 | ): |
| 391 | print(f"Exporting Chunk {chunk_idx} to PTE") |
| 392 | example_inputs, dynamic_shapes = model.get_example_inputs( |
| 393 | max_num_token, max_cache_size, True |
| 394 | ) |
| 395 | print("Getting pre autograd ATen Dialect Graph") |
| 396 | pre_autograd_aten_dialect = torch.export.export( |
| 397 | model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True |
| 398 | ).module() # NOTE: Will be replaced with export |
| 399 | quantizer = NeuropilotQuantizer() |
| 400 | quantizer.setup_precision(getattr(Precision, precision)) |
| 401 | prepared_graph = prepare_pt2e(pre_autograd_aten_dialect, quantizer) |
| 402 | # at this point quant min max are inf |
| 403 | if cal_dataset is not None: |
| 404 | calibrate_model(prepared_graph, cal_dataset, str(chunk_idx)) |
| 405 | else: |
| 406 | prepared_graph(*example_inputs) # dummy calibration |
| 407 | converted_graph = convert_pt2e(prepared_graph, fold_quantize=False) |
| 408 | |
| 409 | method_to_edge_program = {} |
| 410 | method_to_partitioner = {} |
| 411 | edge_compile_config = exir.EdgeCompileConfig(_check_ir_validity=False) |
| 412 | |
| 413 | model_shared_key_name = f"{exp_name}_{chunk_idx}" |
| 414 | |
| 415 | # Fixed Shape Export Here |
| 416 | for shape, ntok_and_cache in export_shapes.items(): |
| 417 | model_fname = f"{exp_name}_{shape}_{chunk_idx}" |
| 418 | example_inputs = model.get_example_inputs(*ntok_and_cache) |
| 419 | print(f"Getting ATen Dialect Graph for {exp_name} {shape} chunk {chunk_idx}") |
| 420 | aten_dialect: exir.ExportedProgram = torch.export.export( |
| 421 | converted_graph, example_inputs, strict=True |
| 422 | ) |
| 423 | |
| 424 | method_to_edge_program[f"{model_fname}"] = exir.to_edge( |
| 425 | aten_dialect |
| 426 | ).exported_program() |
| 427 | del aten_dialect |
| 428 | |
| 429 | compile_spec = [ |
| 430 | CompileSpec("gno", b"LTS"), |
| 431 | CompileSpec("gno-exp", b""), |
| 432 | CompileSpec("gno-non-4d-tiling", b""), |
| 433 | CompileSpec("ImportForever", struct.pack("?", True)), |
| 434 | CompileSpec("platform-config", platform_b), |
| 435 | CompileSpec("ExtractSharedBlobKey", model_shared_key_name.encode()), |
| 436 | ] |
no test coverage detected