Commit 6634a0e0 authored by zhuwenwen's avatar zhuwenwen
Browse files

support llama model tn/nn

parent a10e9cee
...@@ -24,7 +24,8 @@ def _set_default_torch_dtype(dtype: torch.dtype): ...@@ -24,7 +24,8 @@ def _set_default_torch_dtype(dtype: torch.dtype):
def _get_model_architecture(model_config: ModelConfig) -> Type[nn.Module]: def _get_model_architecture(model_config: ModelConfig) -> Type[nn.Module]:
architectures = getattr(model_config.hf_config, "architectures", []) architectures = getattr(model_config.hf_config, "architectures", [])
if architectures == ['LlamaForCausalLM']: if architectures == ['LlamaForCausalLM']:
os.environ['LLAMA_NN'] = '1' if os.getenv('LLAMA_NN') != '0':
os.environ['LLAMA_NN'] = '1'
# Special handling for quantized Mixtral. # Special handling for quantized Mixtral.
# FIXME(woosuk): This is a temporary hack. # FIXME(woosuk): This is a temporary hack.
if (model_config.quantization is not None if (model_config.quantization is not None
......
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