(
output_folder,
exp_name,
model,
precision,
max_num_token,
max_cache_size,
chunk_idx,
export_shapes,
platform_b,
cal_dataset=None,
)
| 394 | |
| 395 | |
| 396 | def export_to_et_ir( |
| 397 | output_folder, |
| 398 | exp_name, |
| 399 | model, |
| 400 | precision, |
| 401 | max_num_token, |
| 402 | max_cache_size, |
| 403 | chunk_idx, |
| 404 | export_shapes, |
| 405 | platform_b, |
| 406 | cal_dataset=None, |
| 407 | ): |
| 408 | print(f"Exporting Chunk {chunk_idx} to PTE") |
| 409 | example_inputs, dynamic_shapes = model.get_example_inputs( |
| 410 | max_num_token, max_cache_size, True |
| 411 | ) |
| 412 | print("Getting pre autograd ATen Dialect Graph") |
| 413 | pre_autograd_aten_dialect = torch.export.export( |
| 414 | model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True |
| 415 | ).module() # NOTE: Will be replaced with export |
| 416 | quantizer = NeuropilotQuantizer() |
| 417 | quantizer.setup_precision(getattr(Precision, precision)) |
| 418 | prepared_graph = prepare_pt2e(pre_autograd_aten_dialect, quantizer) |
| 419 | # at this point quant min max are inf |
| 420 | if cal_dataset is not None: |
| 421 | calibrate_model(prepared_graph, cal_dataset, str(chunk_idx)) |
| 422 | else: |
| 423 | prepared_graph(*example_inputs) # dummy calibration |
| 424 | converted_graph = convert_pt2e(prepared_graph, fold_quantize=False) |
| 425 | |
| 426 | print("Getting ATen Dialect Graph") |
| 427 | # Fixed Shape Export Here |
| 428 | for shape, ntok_and_cache in export_shapes.items(): |
| 429 | dest_path = get_dest_path(output_folder, exp_name, shape, chunk_idx) |
| 430 | print(f"Exporting Shape {shape} to:\n{dest_path}") |
| 431 | example_inputs = model.get_example_inputs(*ntok_and_cache) |
| 432 | aten_dialect: exir.ExportedProgram = torch.export.export( |
| 433 | converted_graph, example_inputs, strict=True |
| 434 | ) |
| 435 | |
| 436 | print("Lowering to Edge Dialect Graph") |
| 437 | edge_program: exir.EdgeProgramManager = exir.to_edge( |
| 438 | aten_dialect, |
| 439 | compile_config=exir.EdgeCompileConfig(_check_ir_validity=False), |
| 440 | ) |
| 441 | del aten_dialect |
| 442 | |
| 443 | print("Delegating Edge Program to Neuropilot Backend") |
| 444 | compile_spec = [ |
| 445 | CompileSpec("gno", b"LTS"), |
| 446 | CompileSpec("gno-exp", b""), |
| 447 | CompileSpec("gno-non-4d-tiling", b""), |
| 448 | CompileSpec("ImportForever", struct.pack("?", True)), |
| 449 | CompileSpec("platform-config", platform_b), |
| 450 | ] |
| 451 | partitioner = NeuropilotPartitioner(compile_spec) |
| 452 | delegated_program = edge_program.to_backend(partitioner) |
| 453 | print("Exported Delegated Program:") |
no test coverage detected