├── .gitignore ├── LICENSE ├── code └── jet_nemotron_models ├── README.md ├── assets ├── hardware-aware.png ├── jet-nemotron.png ├── jetblock.png ├── main-results.png ├── postnas-roadmap.png └── search-results.png ├── jetai ├── evaluation │ ├── configs │ │ ├── bbh.yaml │ │ ├── commonsense.yaml │ │ ├── cruxeval.yaml │ │ ├── evalplus.yaml │ │ ├── longbench.yaml │ │ ├── math_gen.yaml │ │ ├── math_mc.yaml │ │ ├── mmlu.yaml │ │ ├── mmlu_pro.yaml │ │ └── retrieval.yaml │ ├── lm_eval_harness │ │ ├── __init__.py │ │ ├── evaluator.py │ │ └── wrapper.py │ └── meta_eval.py ├── inference │ ├── generate.py │ └── measure_throuput.py ├── modeling │ ├── __init__.py │ ├── configs │ │ ├── jetlm_2B.json │ │ └── jetlm_4B.json │ ├── hf │ │ ├── added_tokens.json │ │ ├── chat_template.jinja │ │ ├── configuration_jet_nemotron.py │ │ ├── dconv_fwd_cache.py │ │ ├── dconv_fwdbwd.py │ │ ├── dconv_step.py │ │ ├── dynamic_conv.py │ │ ├── generation_config.json │ │ ├── jet_block.py │ │ ├── kv_cache.py │ │ ├── merges.txt │ │ ├── modeling_jet_nemotron.py │ │ ├── special_tokens_map.json │ │ ├── tokenizer.json │ │ ├── tokenizer_config.json │ │ └── vocab.json │ └── wrapper.py └── utils │ ├── __init__.py │ ├── config.py │ ├── debug.py │ ├── dist.py │ └── misc.py ├── pyproject.toml ├── requirements.txt └── scripts └── eval ├── 2B ├── bbh.sh ├── code.sh ├── commonsense.sh ├── full.sh ├── longbench.sh ├── math.sh ├── mmlu.sh ├── mmlu_pro.sh └── retrieval.sh └── 4B ├── bbh.sh ├── code.sh ├── commonsense.sh ├── full.sh ├── longbench.sh ├── math.sh ├── mmlu.sh ├── mmlu_pro.sh └── retrieval.sh /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE/code: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/LICENSE/code -------------------------------------------------------------------------------- /LICENSE/jet_nemotron_models: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/LICENSE/jet_nemotron_models -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/README.md -------------------------------------------------------------------------------- /assets/hardware-aware.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/assets/hardware-aware.png -------------------------------------------------------------------------------- /assets/jet-nemotron.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/assets/jet-nemotron.png -------------------------------------------------------------------------------- /assets/jetblock.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/assets/jetblock.png -------------------------------------------------------------------------------- /assets/main-results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/assets/main-results.png -------------------------------------------------------------------------------- /assets/postnas-roadmap.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/assets/postnas-roadmap.png -------------------------------------------------------------------------------- /assets/search-results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/assets/search-results.png -------------------------------------------------------------------------------- /jetai/evaluation/configs/bbh.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/configs/bbh.yaml -------------------------------------------------------------------------------- /jetai/evaluation/configs/commonsense.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/configs/commonsense.yaml -------------------------------------------------------------------------------- /jetai/evaluation/configs/cruxeval.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/configs/cruxeval.yaml -------------------------------------------------------------------------------- /jetai/evaluation/configs/evalplus.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/configs/evalplus.yaml -------------------------------------------------------------------------------- /jetai/evaluation/configs/longbench.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/configs/longbench.yaml -------------------------------------------------------------------------------- /jetai/evaluation/configs/math_gen.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/configs/math_gen.yaml -------------------------------------------------------------------------------- /jetai/evaluation/configs/math_mc.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/configs/math_mc.yaml -------------------------------------------------------------------------------- /jetai/evaluation/configs/mmlu.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/configs/mmlu.yaml -------------------------------------------------------------------------------- /jetai/evaluation/configs/mmlu_pro.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/configs/mmlu_pro.yaml -------------------------------------------------------------------------------- /jetai/evaluation/configs/retrieval.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/configs/retrieval.yaml -------------------------------------------------------------------------------- /jetai/evaluation/lm_eval_harness/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/lm_eval_harness/__init__.py -------------------------------------------------------------------------------- /jetai/evaluation/lm_eval_harness/evaluator.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/lm_eval_harness/evaluator.py -------------------------------------------------------------------------------- /jetai/evaluation/lm_eval_harness/wrapper.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/lm_eval_harness/wrapper.py -------------------------------------------------------------------------------- /jetai/evaluation/meta_eval.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/evaluation/meta_eval.py -------------------------------------------------------------------------------- /jetai/inference/generate.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/inference/generate.py -------------------------------------------------------------------------------- /jetai/inference/measure_throuput.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/inference/measure_throuput.py -------------------------------------------------------------------------------- /jetai/modeling/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/__init__.py -------------------------------------------------------------------------------- /jetai/modeling/configs/jetlm_2B.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/configs/jetlm_2B.json -------------------------------------------------------------------------------- /jetai/modeling/configs/jetlm_4B.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/configs/jetlm_4B.json -------------------------------------------------------------------------------- /jetai/modeling/hf/added_tokens.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/added_tokens.json -------------------------------------------------------------------------------- /jetai/modeling/hf/chat_template.jinja: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/chat_template.jinja -------------------------------------------------------------------------------- /jetai/modeling/hf/configuration_jet_nemotron.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/configuration_jet_nemotron.py -------------------------------------------------------------------------------- /jetai/modeling/hf/dconv_fwd_cache.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/dconv_fwd_cache.py -------------------------------------------------------------------------------- /jetai/modeling/hf/dconv_fwdbwd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/dconv_fwdbwd.py -------------------------------------------------------------------------------- /jetai/modeling/hf/dconv_step.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/dconv_step.py -------------------------------------------------------------------------------- /jetai/modeling/hf/dynamic_conv.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/dynamic_conv.py -------------------------------------------------------------------------------- /jetai/modeling/hf/generation_config.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/generation_config.json -------------------------------------------------------------------------------- /jetai/modeling/hf/jet_block.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/jet_block.py -------------------------------------------------------------------------------- /jetai/modeling/hf/kv_cache.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/kv_cache.py -------------------------------------------------------------------------------- /jetai/modeling/hf/merges.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/merges.txt -------------------------------------------------------------------------------- /jetai/modeling/hf/modeling_jet_nemotron.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/modeling_jet_nemotron.py -------------------------------------------------------------------------------- /jetai/modeling/hf/special_tokens_map.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/special_tokens_map.json -------------------------------------------------------------------------------- /jetai/modeling/hf/tokenizer.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/tokenizer.json -------------------------------------------------------------------------------- /jetai/modeling/hf/tokenizer_config.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/tokenizer_config.json -------------------------------------------------------------------------------- /jetai/modeling/hf/vocab.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/hf/vocab.json -------------------------------------------------------------------------------- /jetai/modeling/wrapper.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/modeling/wrapper.py -------------------------------------------------------------------------------- /jetai/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/utils/__init__.py -------------------------------------------------------------------------------- /jetai/utils/config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/utils/config.py -------------------------------------------------------------------------------- /jetai/utils/debug.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/utils/debug.py -------------------------------------------------------------------------------- /jetai/utils/dist.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/utils/dist.py -------------------------------------------------------------------------------- /jetai/utils/misc.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/jetai/utils/misc.py -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/pyproject.toml -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/requirements.txt -------------------------------------------------------------------------------- /scripts/eval/2B/bbh.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/2B/bbh.sh -------------------------------------------------------------------------------- /scripts/eval/2B/code.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/2B/code.sh -------------------------------------------------------------------------------- /scripts/eval/2B/commonsense.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/2B/commonsense.sh -------------------------------------------------------------------------------- /scripts/eval/2B/full.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/2B/full.sh -------------------------------------------------------------------------------- /scripts/eval/2B/longbench.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/2B/longbench.sh -------------------------------------------------------------------------------- /scripts/eval/2B/math.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/2B/math.sh -------------------------------------------------------------------------------- /scripts/eval/2B/mmlu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/2B/mmlu.sh -------------------------------------------------------------------------------- /scripts/eval/2B/mmlu_pro.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/2B/mmlu_pro.sh -------------------------------------------------------------------------------- /scripts/eval/2B/retrieval.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/2B/retrieval.sh -------------------------------------------------------------------------------- /scripts/eval/4B/bbh.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/4B/bbh.sh -------------------------------------------------------------------------------- /scripts/eval/4B/code.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/4B/code.sh -------------------------------------------------------------------------------- /scripts/eval/4B/commonsense.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/4B/commonsense.sh -------------------------------------------------------------------------------- /scripts/eval/4B/full.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/4B/full.sh -------------------------------------------------------------------------------- /scripts/eval/4B/longbench.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/4B/longbench.sh -------------------------------------------------------------------------------- /scripts/eval/4B/math.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/4B/math.sh -------------------------------------------------------------------------------- /scripts/eval/4B/mmlu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/4B/mmlu.sh -------------------------------------------------------------------------------- /scripts/eval/4B/mmlu_pro.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/4B/mmlu_pro.sh -------------------------------------------------------------------------------- /scripts/eval/4B/retrieval.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/Jet-Nemotron/HEAD/scripts/eval/4B/retrieval.sh --------------------------------------------------------------------------------