From f9d954475efaf2187f12c5fc79b1913a72f23542 Mon Sep 17 00:00:00 2001 From: Scott Wolchok Date: Tue, 9 Apr 2024 10:03:13 -0700 Subject: [PATCH] generate.py: Make DSO path absolute Summary: Relative path doesn't seem to work with dlopen. Test Plan: generate.py with --dso-path was failing to dlopen before; now it works. --- generate.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/generate.py b/generate.py index eb7a809ae..df4ae4b7b 100644 --- a/generate.py +++ b/generate.py @@ -327,7 +327,7 @@ def main( if (checkpoint_path and (dso_path or pte_path)): print("Warning: checkpoint path ignored because an exported DSO or PTE path specified") - + if not tokenizer_path: assert checkpoint_path, "either a tokenizer or a checkpoint path must be specified" tokenizer_path = checkpoint_path.parent / "tokenizer.model" @@ -359,7 +359,7 @@ def main( # attributes will NOT be seen on by AOTI-compiled forward # function, e.g. calling model.setup_cache will NOT touch # AOTI compiled and maintained model buffers such as kv_cache. - model.forward = torch._export.aot_load(str(dso_path), device) + model.forward = torch._export.aot_load(str(dso_path.absolute()), device) except: raise RuntimeError(f"Failed to load AOTI compiled {dso_path}") elif pte_path: @@ -596,7 +596,7 @@ def cli(): if args.seed: torch.manual_seed(args.seed) - + main( args.prompt, args.interactive,