Support spec decoding when LoRA is applied to target model (#12903)
This commit is contained in:
Lifu Huang
@@ -395,3 +395,177 @@ def ensure_reproducibility():
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.use_deterministic_algorithms(True)
|
||||
|
||||
|
||||
TEST_MULTIPLE_BATCH_PROMPTS = [
|
||||
"""
|
||||
### Instruction:
|
||||
Tell me about llamas and alpacas
|
||||
### Response:
|
||||
Llamas are large, long-necked animals with a woolly coat. They have two toes on each foot instead of three like other camelids (camels, dromedaries). Llamas live in the Andean mountains of South America where they graze on grasses and shrubs. Alpaca is another name for domesticated llama. The word "alpaca" comes from an Incan language meaning "golden fleece." Alpacas look very similar to llamas but are smaller than their wild relatives. Both species were used by ancient people as pack animals and for meat. Today both llamas and alpacas are raised primarily for their fiber which can be spun into yarn or knitted into clothing.
|
||||
### Question 2:
|
||||
What do you know about llamas?
|
||||
### Answer:
|
||||
""",
|
||||
"""
|
||||
### Instruction:
|
||||
Write a poem about the transformers Python library.
|
||||
Mention the word "large language models" in that poem.
|
||||
### Response:
|
||||
The Transformers are large language models,
|
||||
They're used to make predictions on text.
|
||||
""",
|
||||
"AI is a field of computer science focused on",
|
||||
"Computer science is the study of",
|
||||
"Write a short story.",
|
||||
"What are the main components of a computer?",
|
||||
]
|
||||
|
||||
|
||||
def create_multiple_batch_test_samples(
|
||||
prompts: List[str], lora_adapter_paths: List[str]
|
||||
):
|
||||
random.seed(42)
|
||||
|
||||
return [
|
||||
(
|
||||
[
|
||||
random.choice(prompts),
|
||||
random.choice(prompts),
|
||||
random.choice(prompts),
|
||||
],
|
||||
[
|
||||
None,
|
||||
lora_adapter_paths[0],
|
||||
lora_adapter_paths[1],
|
||||
],
|
||||
),
|
||||
(
|
||||
[
|
||||
random.choice(prompts),
|
||||
random.choice(prompts),
|
||||
random.choice(prompts),
|
||||
],
|
||||
[
|
||||
lora_adapter_paths[0],
|
||||
None,
|
||||
lora_adapter_paths[1],
|
||||
],
|
||||
),
|
||||
(
|
||||
[
|
||||
random.choice(prompts),
|
||||
random.choice(prompts),
|
||||
random.choice(prompts),
|
||||
],
|
||||
[lora_adapter_paths[0], lora_adapter_paths[1], None],
|
||||
),
|
||||
(
|
||||
[
|
||||
random.choice(prompts),
|
||||
random.choice(prompts),
|
||||
random.choice(prompts),
|
||||
],
|
||||
[None, lora_adapter_paths[1], None],
|
||||
),
|
||||
(
|
||||
[
|
||||
random.choice(prompts),
|
||||
random.choice(prompts),
|
||||
random.choice(prompts),
|
||||
],
|
||||
[None, None, None],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def run_lora_multiple_batch_on_model_cases(
|
||||
model_cases: List[LoRAModelCase],
|
||||
use_spec_decoding: bool = False,
|
||||
attention_backend: str = "torch_native",
|
||||
disable_cuda_graph: bool = True,
|
||||
enable_deterministic_inference: bool = False,
|
||||
):
|
||||
for model_case in model_cases:
|
||||
for torch_dtype in TORCH_DTYPES:
|
||||
max_new_tokens = 32
|
||||
base_path = model_case.base
|
||||
lora_adapter_paths = [a.name for a in model_case.adaptors]
|
||||
assert len(lora_adapter_paths) >= 2
|
||||
|
||||
batches = create_multiple_batch_test_samples(
|
||||
TEST_MULTIPLE_BATCH_PROMPTS, lora_adapter_paths
|
||||
)
|
||||
|
||||
print(
|
||||
f"\n========== Testing multiple batches on base '{base_path}', dtype={torch_dtype} ---"
|
||||
)
|
||||
|
||||
# Initialize runners
|
||||
ensure_reproducibility()
|
||||
spec_args = (
|
||||
{}
|
||||
if not use_spec_decoding
|
||||
else {
|
||||
"speculative_algorithm": "NGRAM",
|
||||
"speculative_num_draft_tokens": 5,
|
||||
"speculative_ngram_min_match_window_size": 2,
|
||||
"speculative_ngram_max_match_window_size": 15,
|
||||
}
|
||||
)
|
||||
srt_runner = SRTRunner(
|
||||
base_path,
|
||||
torch_dtype=torch_dtype,
|
||||
model_type="generation",
|
||||
lora_paths=[lora_adapter_paths[0], lora_adapter_paths[1]],
|
||||
max_loras_per_batch=len(lora_adapter_paths) + 1,
|
||||
sleep_on_idle=True, # Eliminate non-determinism by forcing all requests to be processed in one batch.
|
||||
attention_backend=attention_backend,
|
||||
enable_deterministic_inference=enable_deterministic_inference,
|
||||
disable_cuda_graph=disable_cuda_graph,
|
||||
**spec_args,
|
||||
)
|
||||
|
||||
ensure_reproducibility()
|
||||
hf_runner = HFRunner(
|
||||
base_path,
|
||||
torch_dtype=torch_dtype,
|
||||
model_type="generation",
|
||||
patch_model_do_sample_false=True,
|
||||
)
|
||||
|
||||
with srt_runner, hf_runner:
|
||||
for i, (prompts, lora_paths) in enumerate(batches):
|
||||
print(
|
||||
f"\n--- Running Batch {i+1} --- prompts: {prompts}, lora_paths: {lora_paths}"
|
||||
)
|
||||
|
||||
srt_outputs = srt_runner.batch_forward(
|
||||
prompts,
|
||||
max_new_tokens=max_new_tokens,
|
||||
lora_paths=lora_paths,
|
||||
)
|
||||
|
||||
hf_outputs = hf_runner.forward(
|
||||
prompts,
|
||||
max_new_tokens=max_new_tokens,
|
||||
lora_paths=lora_paths,
|
||||
)
|
||||
|
||||
print("SRT outputs:", [s for s in srt_outputs.output_strs])
|
||||
print("HF outputs:", [s for s in hf_outputs.output_strs])
|
||||
|
||||
for srt_out, hf_out in zip(
|
||||
srt_outputs.output_strs, hf_outputs.output_strs
|
||||
):
|
||||
srt_str = srt_out.strip()
|
||||
hf_str = hf_out.strip()
|
||||
rouge_tol = model_case.rouge_l_tolerance
|
||||
rouge_score = calculate_rouge_l([srt_str], [hf_str])[0]
|
||||
if rouge_score < rouge_tol:
|
||||
raise AssertionError(
|
||||
f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} "
|
||||
f"for base '{base_path}', adaptor '{lora_paths}', prompt: '{prompts}...'"
|
||||
)
|
||||
|
||||
print(f"--- Batch {i+1} Comparison Passed --- ")
|
||||
|
||||
Reference in New Issue
Block a user