Add default 'auto' MODEL_IMPL_TYPE that resolves based on architecture

[xingliu14]

Description

• Add auto as default value for MODEL_IMPL_TYPE env var
• For GptOssForCausalLM, auto resolves to vllm for better performance
• For all other architectures, 'auto' resolves to flax_nnx for better performance
• Add tests for 'auto' resolution behavior

Tests

pytest

tests/test_envs.py
tests/models/common/test_model_loader.py

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[xingliu14]

@kyuyeunk Please review.

[kyuyeunk]

Wouldn't it be possible move 'auto' to 'match/case' as well?

[xingliu14]

It is possible to move it in to match-case, but in that case it will have duplicated codes, including: get_vllm_model, get_flax_model and the fall back check. I think resolve first then use the same code will be more clean.

[kyuyeunk]

I don't see any update? Can you verify if the changes were pushed?

[kyuyeunk]

LGTM. Thank you for working on it!

[kyuyeunk]

seems like CI is failing. Have you rebase the branch to HEAD?

[xingliu14]

I believe so, let me look into it.

[xingliu14]

Looks like tests/layers/vllm/test_awq.py is failing from the main branch as well.

[kyuyeunk]

seems like it's due to an upstream change. let me create a quick fix for this.

[kyuyeunk]

Please wait until this is merged: #1284

[kyuyeunk]

the PR has been merged. please update the branch and try again.

[kyuyeunk]

https://github.com/vllm-project/tpu-inference/actions/runs/20125295806/job/57753569146?pr=1255

pre-commit hook(s) made changes.
If you are seeing this message in CI, reproduce locally with: `pre-commit run --all-files`.
To run `pre-commit` as part of git workflow, use `pre-commit install`.
All changes made by hooks:
diff --git a/tpu_inference/models/common/model_loader.py b/tpu_inference/models/common/model_loader.py
index f260e5f..fb035bd 100644
--- a/tpu_inference/models/common/model_loader.py
+++ b/tpu_inference/models/common/model_loader.py
@@ -27,7 +27,8 @@ _MODEL_REGISTRY = {}
 # Architectures that prefer "vllm" implementation type when MODEL_IMPL_TYPE is "auto".
 # These architectures are listed here because they have better performance with the
 # vLLM PyTorch backend compared to the flax_nnx JAX backend for now.
-_VLLM_PREFERRED_ARCHITECTURES: frozenset[str] = frozenset({"GptOssForCausalLM"})
+_VLLM_PREFERRED_ARCHITECTURES: frozenset[str] = frozenset(
+    {"GptOssForCausalLM"})
 
 
 class UnsupportedArchitectureError(ValueError):
Error: Process completed with exit code 1.

Please fix pre-commit failure.

[kyuyeunk]

Thank you so much for making this feature!