air_llm/tests/test_streaming_gpu.py is a manual GPU test harness for AirLLM layer-streaming inference; it is not an automated pytest suite and must be invoked explicitly.[1]
The harness can cap the visible VRAM for the test process via torch.cuda.set_per_process_memory_fraction, emulating a smaller GPU card without physically needing one.[1] When the --max-vram-gb value equals or exceeds the device's total memory, the harness skips capping and prints a warning instead of applying the fraction.[1]
The harness invokes AirLLM via AutoModel.from_pretrained with compression and delete_original forwarded from CLI flags, and calls model.generate with do_sample=False (greedy decoding) and return_dict_in_generate=True.[1] --compression accepts only None, "4bit", or "8bit", corresponding to the compression options exposed by AutoModel.from_pretrained — see Compression for details on those modes.[1] The default prompt is "The capital of France is" and the default number of new tokens is 12.[1]
Peak VRAM usage (in MB) is reported via torch.cuda.max_memory_allocated(), with the peak counter reset immediately before the model.generate() call so only inference is measured.[1]
When --compare is passed, the harness runs the full-load Transformers reference before capping VRAM so the complete model fits on the GPU; only then is cap_vram applied for the AirLLM run.[1] The reference full-load run loads the model with dtype=torch.float16 and is documented as only feasible for small models.[1] After the reference run completes, the harness explicitly deletes the reference model and calls torch.cuda.empty_cache() to free VRAM before the AirLLM run starts.[1] --compare mode asserts token-level output equality between the AirLLM streaming run and the full-load Transformers run, exiting with code 1 on mismatch.[1]
How to run the GPU streaming test for a tiny model and verify output matches a full-load reference:
python test_streaming_gpu.py --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --compare
How to emulate a 4 GB GPU card running a 7B model with air_llm/tests/test_streaming_gpu.py:
python test_streaming_gpu.py --model Qwen/Qwen2.5-7B-Instruct --max-vram-gb 4
The canonical usage of AirLLMLlama2 in air_llm/inference_example.py accepts either a Hugging Face repo ID or a local model path as the first argument:
model = AirLLMLlama2("garage-bAInd/Platypus2-70B-instruct")
# or
model = AirLLMLlama2("/home/ubuntu/.cache/huggingface/hub/...")
The end-to-end inference pattern in air_llm/inference_example.py tokenizes with return_attention_mask=False, passes input_ids to .cuda(), calls .generate() with use_cache=True and return_dict_in_generate=True, then decodes generation_output.sequences[0]:
input_tokens = model.tokenizer(input_text, return_tensors="pt",
return_attention_mask=False, truncation=True,
max_length=MAX_LENGTH, padding=True)
generation_output = model.generate(
input_tokens['input_ids'].cuda(),
max_new_tokens=2, use_cache=True, return_dict_in_generate=True)
output = model.tokenizer.decode(generation_output.sequences[0])
Sources