Ubuntu

tests

5ee43e9 9 months ago

2.31 kB

#!/usr/bin/env python3
# T5 decoder (no cache) on Neuron – constant shapes, full graph, no Apex
import os
os.environ["USE_FUSED_LAYER_NORM"] = "0"  # MUST be before any transformers import

import argparse
import logging
import time
import torch
from transformers import T5Tokenizer, T5Model
import torch_neuronx  # guarantees Neuron backend

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def main():
    parser = argparse.ArgumentParser(description="T5 decoder on Neuron (full graph, no cache)")
    parser.add_argument("--model", default="t5-small")
    parser.add_argument("--seq-len", type=int, default=128, help="Fixed seq length")
    args = parser.parse_args()

    torch.manual_seed(42)
    torch.set_default_dtype(torch.float32)

    tokenizer = T5Tokenizer.from_pretrained(args.model)
    # disable DynamicCache → no deepcopy of config
    model = T5Model.from_pretrained(
        args.model,
        torch_dtype=torch.float32,
        attn_implementation="eager",
        use_cache=False,  # <-- static shapes, no cache
    ).eval()

    # constant-shape inputs
    text = "hello"
    enc_tok = tokenizer(text, max_length=args.seq_len, padding="max_length", truncation=True, return_tensors="pt")
    with torch.no_grad():
        enc_out = model.encoder(input_ids=enc_tok.input_ids).last_hidden_state.detach()

    dec_tok = tokenizer("<pad>", max_length=args.seq_len, padding="max_length", return_tensors="pt")

    # pre-run to lock shapes
    with torch.no_grad():
        _ = model.decoder(input_ids=dec_tok.input_ids, encoder_hidden_states=enc_out).last_hidden_state

    # compile decoder forward only (full graph)
    decode_fn = lambda inp, enc: model.decoder(input_ids=inp, encoder_hidden_states=enc).last_hidden_state
    decode_fn = torch.compile(decode_fn, backend="neuron", fullgraph=True)

    # warmup
    start = time.time()
    with torch.no_grad():
        _ = decode_fn(dec_tok.input_ids, enc_out)
    logger.info("Warmup: %.3f s", time.time() - start)

    # benchmark
    start = time.time()
    with torch.no_grad():
        hidden = decode_fn(dec_tok.input_ids, enc_out)
    logger.info("Run: %.3f s", time.time() - start)
    logger.info("Hidden shape: %s", hidden.shape)  # [B, seq_len, d_model]


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	# T5 decoder (no cache) on Neuron – constant shapes, full graph, no Apex
	import os
	os.environ["USE_FUSED_LAYER_NORM"] = "0" # MUST be before any transformers import

	import argparse
	import logging
	import time
	import torch
	from transformers import T5Tokenizer, T5Model
	import torch_neuronx # guarantees Neuron backend

	logging.basicConfig(level=logging.INFO)
	logger = logging.getLogger(__name__)


	def main():
	parser = argparse.ArgumentParser(description="T5 decoder on Neuron (full graph, no cache)")
	parser.add_argument("--model", default="t5-small")
	parser.add_argument("--seq-len", type=int, default=128, help="Fixed seq length")
	args = parser.parse_args()

	torch.manual_seed(42)
	torch.set_default_dtype(torch.float32)

	tokenizer = T5Tokenizer.from_pretrained(args.model)
	# disable DynamicCache → no deepcopy of config
	model = T5Model.from_pretrained(
	args.model,
	torch_dtype=torch.float32,
	attn_implementation="eager",
	use_cache=False, # <-- static shapes, no cache
	).eval()

	# constant-shape inputs
	text = "hello"
	enc_tok = tokenizer(text, max_length=args.seq_len, padding="max_length", truncation=True, return_tensors="pt")
	with torch.no_grad():
	enc_out = model.encoder(input_ids=enc_tok.input_ids).last_hidden_state.detach()

	dec_tok = tokenizer("<pad>", max_length=args.seq_len, padding="max_length", return_tensors="pt")

	# pre-run to lock shapes
	with torch.no_grad():
	_ = model.decoder(input_ids=dec_tok.input_ids, encoder_hidden_states=enc_out).last_hidden_state

	# compile decoder forward only (full graph)
	decode_fn = lambda inp, enc: model.decoder(input_ids=inp, encoder_hidden_states=enc).last_hidden_state
	decode_fn = torch.compile(decode_fn, backend="neuron", fullgraph=True)

	# warmup
	start = time.time()
	with torch.no_grad():
	_ = decode_fn(dec_tok.input_ids, enc_out)
	logger.info("Warmup: %.3f s", time.time() - start)

	# benchmark
	start = time.time()
	with torch.no_grad():
	hidden = decode_fn(dec_tok.input_ids, enc_out)
	logger.info("Run: %.3f s", time.time() - start)
	logger.info("Hidden shape: %s", hidden.shape) # [B, seq_len, d_model]


	if __name__ == "__main__":
	main()