File size: 2,869 Bytes
5ee43e9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
#!/usr/bin/env python3
# Phi (Phi-2 default) forward-trace + manual greedy on Neuron – fixed pad token
import argparse
import logging
import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch_neuronx # guarantees Neuron backend
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
@torch.no_grad()
def greedy_generate(model_forward, tokenizer, input_ids, max_new_tokens):
"""Manual greedy loop. Calls the *compiled* forward iteratively."""
B, seq_len = input_ids.shape
device = input_ids.device
position_ids = torch.arange(seq_len, dtype=torch.long, device=device).unsqueeze(0).expand(B, -1)
for _ in range(max_new_tokens):
logits = model_forward(input_ids, position_ids)[0] # unpack tuple
next_id = logits[:, -1, :].argmax(dim=-1, keepdim=True)
input_ids = torch.cat([input_ids, next_id], dim=1)[:, -seq_len:] # rolling window
return input_ids
def main():
parser = argparse.ArgumentParser(description="Phi forward-compile + manual greedy on Neuron")
parser.add_argument("--model", default="microsoft/phi-2")
parser.add_argument("--seq-len", type=int, default=128, help="Fixed context length")
parser.add_argument("--new-tokens", type=int, default=20, help="Tokens to generate")
args = parser.parse_args()
torch.manual_seed(42)
torch.set_default_dtype(torch.float32)
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
# Phi has no pad_token by default
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
args.model,
torch_dtype=torch.float32,
attn_implementation="eager",
use_cache=False, # static shapes
).eval()
prompt = "The future of AI is"
inputs = tokenizer(prompt, max_length=args.seq_len, padding="max_length", truncation=True, return_tensors="pt")
input_ids = inputs.input_ids
B, seq_len = input_ids.shape
# shape lock & compile forward only (full graph)
with torch.no_grad():
position_ids = torch.arange(seq_len, dtype=torch.long).unsqueeze(0).expand(B, -1)
_ = model(input_ids, position_ids)
model.forward = torch.compile(model.forward, backend="neuron", fullgraph=True)
# warmup
start = time.time()
with torch.no_grad():
_ = model(input_ids, position_ids)
logger.info("Warmup (forward): %.3f s", time.time() - start)
# manual greedy generation
start = time.time()
final_ids = greedy_generate(model.forward, tokenizer, input_ids, args.new_tokens)
logger.info("Generate (manual loop): %.3f s", time.time() - start)
text = tokenizer.decode(final_ids[0], skip_special_tokens=True)
logger.info("Output: %s", text)
if __name__ == "__main__":
main() |