Commit 7e4ae96f authored by Your Name's avatar Your Name

Fix --flash-attn CLI flag not being passed to model loading

- Add flash_attn extraction from global_args in _load_default_model()
- Add flash_attn extraction from global_args in _load_model_by_name()
- Now --flash-attn flag will properly enable Flash Attention 2 when loading models
parent 49c01211
...@@ -528,6 +528,8 @@ class MultiModelManager: ...@@ -528,6 +528,8 @@ class MultiModelManager:
kwargs['offload_dir'] = global_args.offload_dir kwargs['offload_dir'] = global_args.offload_dir
if hasattr(global_args, 'ram'): if hasattr(global_args, 'ram'):
kwargs['ram'] = global_args.ram kwargs['ram'] = global_args.ram
if hasattr(global_args, 'flash_attn'):
kwargs['flash_attn'] = global_args.flash_attn
print(f"Loading default model on demand: {self.default_model}") print(f"Loading default model on demand: {self.default_model}")
model_manager.load_model(self.default_model, backend_type=backend_type, **kwargs) model_manager.load_model(self.default_model, backend_type=backend_type, **kwargs)
...@@ -575,6 +577,8 @@ class MultiModelManager: ...@@ -575,6 +577,8 @@ class MultiModelManager:
kwargs['offload_dir'] = global_args.offload_dir kwargs['offload_dir'] = global_args.offload_dir
if hasattr(global_args, 'ram'): if hasattr(global_args, 'ram'):
kwargs['ram'] = global_args.ram kwargs['ram'] = global_args.ram
if hasattr(global_args, 'flash_attn'):
kwargs['flash_attn'] = global_args.flash_attn
print(f"Loading model on demand: {model_name}") print(f"Loading model on demand: {model_name}")
model_manager.load_model(model_name, backend_type=backend_type, **kwargs) model_manager.load_model(model_name, backend_type=backend_type, **kwargs)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment