(
output_folder,
exp_name,
model,
precision,
max_num_token,
max_cache_size,
chunk_idx,
export_shapes,
platform_b,
cal_dataset=None,
)
| 323 | |
| 324 | |
| 325 | def export_to_et_ir( |
| 326 | output_folder, |
| 327 | exp_name, |
| 328 | model, |
| 329 | precision, |
| 330 | max_num_token, |
| 331 | max_cache_size, |
| 332 | chunk_idx, |
| 333 | export_shapes, |
| 334 | platform_b, |
| 335 | cal_dataset=None, |
| 336 | ): |
| 337 | print(f"Exporting Chunk {chunk_idx} to PTE") |
| 338 | example_inputs, dynamic_shapes = model.get_example_inputs( |
| 339 | max_num_token, max_cache_size, True |
| 340 | ) |
| 341 | print("Getting pre autograd ATen Dialect Graph") |
| 342 | pre_autograd_aten_dialect = torch.export.export( |
| 343 | model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True |
| 344 | ).module() # NOTE: Will be replaced with export |
| 345 | quantizer = NeuropilotQuantizer() |
| 346 | quantizer.setup_precision(getattr(Precision, precision)) |
| 347 | prepared_graph = prepare_pt2e(pre_autograd_aten_dialect, quantizer) |
| 348 | # at this point quant min max are inf |
| 349 | if cal_dataset is not None: |
| 350 | calibrate_model(prepared_graph, cal_dataset, str(chunk_idx)) |
| 351 | else: |
| 352 | prepared_graph(*example_inputs) # dummy calibration |
| 353 | converted_graph = convert_pt2e(prepared_graph, fold_quantize=False) |
| 354 | |
| 355 | method_to_edge_program = {} |
| 356 | method_to_partitioner = {} |
| 357 | edge_compile_config = exir.EdgeCompileConfig(_check_ir_validity=False) |
| 358 | |
| 359 | model_shared_key_name = f"{exp_name}_{chunk_idx}" |
| 360 | |
| 361 | # Fixed Shape Export Here |
| 362 | for shape, ntok_and_cache in export_shapes.items(): |
| 363 | model_fname = f"{exp_name}_{shape}_{chunk_idx}" |
| 364 | example_inputs = model.get_example_inputs(*ntok_and_cache) |
| 365 | print(f"Getting ATen Dialect Graph for {exp_name} {shape} chunk {chunk_idx}") |
| 366 | aten_dialect: exir.ExportedProgram = torch.export.export( |
| 367 | converted_graph, example_inputs, strict=True |
| 368 | ) |
| 369 | |
| 370 | method_to_edge_program[f"{model_fname}"] = exir.to_edge( |
| 371 | aten_dialect |
| 372 | ).exported_program() |
| 373 | del aten_dialect |
| 374 | |
| 375 | compile_spec = [ |
| 376 | CompileSpec("gno", b"LTS"), |
| 377 | CompileSpec("gno-exp", b""), |
| 378 | CompileSpec("gno-non-4d-tiling", b""), |
| 379 | CompileSpec("ImportForever", struct.pack("?", True)), |
| 380 | CompileSpec("platform-config", platform_b), |
| 381 | CompileSpec("ExtractSharedBlobKey", model_shared_key_name.encode()), |
| 382 | ] |
no test coverage detected