INTERNAL ASSERT FAILED XLANativeFunctions.cpp :596

username: tiagowright
cooked: <p>I’m training a simple Gpt model on AWS Trn instance, and getting the error:</p>
<pre><code class="lang-auto">RuntimeError: common_device INTERNAL ASSERT FAILED at "bazel-out/k8-opt/bin/torch_xla/csrc/XLANativeFunctions.cpp":596, please report a bug to PyTorch.
</code></pre>
<p>The code runs on the CPU, but fails when the device is the accelerator. The offending line seems to be this one, and the runtime error seems to go away when it is commented out:</p>
<pre><code class="lang-auto">wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
</code></pre>
<p>And here is the full script:</p>
<pre><code class="lang-auto">import time
import os
from datetime import datetime
import argparse

import torch
import torch.nn as nn
from torch.nn import functional as F

import torch_xla.core.xla_model as xm


# hyperparameters -- Test
batch_size = 128 # how many independent sequences will we process in parallel?
n_context = 256 # what is the maximum context length for predictions?
max_iters = 5000
eval_interval = 100
learning_rate = 1e-3
eval_iters = 200
n_embd = 128
n_heads = 4
n_layers = 4
dropout = 0.2


corpusfilename = 'tinyshakespeare.txt'


device = 'cpu'
if torch.cuda.is_available():
	device = 'cuda'
elif torch.backends.mps.is_available():
	device = 'mps'
elif xm.xla_device(): ## need to check if this is the right approach
	device = 'xla'

print(f'Running on {device}')


# 
# Start
# 

# read corpus in to inspect it
with open('tinyshakespeare.txt', 'r', encoding='utf-8') as f:
	text = f.read()

# here are all the unique characters that occur in this text
chars = sorted(list(set(text)))
vocab_size = len(chars)

print(f"input corpus has {len(text):,} characters with vocabulary size of {vocab_size} diff characters: {repr(chars)}")

# create a mapping from characters to integers
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }
encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers
decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string

# let's now encode the entire text dataset and store it into a torch.Tensor
data = torch.tensor(encode(text), dtype=torch.long)

# Let's now split up the data into train and validation sets
n = int(0.9*len(data)) # first 90% will be train, rest val
train_data = data[:n]
val_data = data[n:]


def get_batch(split='train'):
	'''
	x, y (B, T)
	'''
	data = train_data if split == 'train' else val_data
	ix = torch.randint(len(data) - n_context, (batch_size,))
	x = torch.stack([data[i:i+n_context] for i in ix])
	y = torch.stack([data[i+1:i+n_context+1] for i in ix])

	return x.to(device), y.to(device)


class Embed(nn.Module):
	
	def __init__(self, n_context, n_vocab, n_embd):
		super().__init__()
		self.token_embedding_table = nn.Embedding(n_vocab, n_embd)
		self.position_embedding_table = nn.Embedding(n_context, n_embd)

	def forward(self, idx):
		B, T = idx.shape

		if T &gt; self.position_embedding_table.weight.shape[0]:
			raise Exception(f'cannot embed inputs that are longer than context window, context window shape {self.position_embedding_table.weight.shape} and input shape {idx.shape}')

		# x is (B, T, E)
		x = self.token_embedding_table(idx)

		# p is (T, E)
		p = self.position_embedding_table(torch.arange(T, device=device))

		# (B, T, E)
		return x + p


class Head(nn.Module):

	def __init__(self, n_context, n_embd, headsize):
		super().__init__()

		self.headsize = headsize
		self.query = nn.Linear(n_embd, headsize)
		self.key = nn.Linear(n_embd, headsize)
		self.value = nn.Linear(n_embd, headsize)
		self.dropout = nn.Dropout(dropout)
		self.register_buffer('tril', torch.tril(torch.ones(n_context, n_context)))

	def forward(self, idx):
		q = self.query(idx)
		k = self.key(idx)
		v = self.value(idx)

		B, T, H = q.shape

		# q, k are (B, T, H)
		# wei is (B, T, T)
		wei = q @ k.transpose(-2, -1) * H**-0.5

		wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
		wei = torch.softmax(wei, dim=-1)
		wei = self.dropout(wei)

		# B, T, T @ B, T, H =&gt; B, T, H
		return wei @ v

class Gpt(nn.Module):
	def __init__(self, n_batch, n_context, n_vocab, n_embd, n_layers, n_heads):
		super().__init__()

		self.n_batch = n_batch
		self.n_context = n_context
		self.n_vocab = n_vocab
		self.n_embd = n_embd
		self.n_layers = n_layers
		self.n_heads = n_heads

		self.embed = Embed(n_context, n_vocab, n_embd)
		self.blocks = Head(n_context, n_embd, n_embd)

	def forward(self, idx, targets):
		'''
		idx (B, T)
		targets (B, n_vocab)
		logits (B, T, n_vocab)
		loss (B*T)
		'''

		x = self.embed(idx)
		x = self.blocks(x)


		if targets is None:
			loss = None

		else:
			B, T, C = logits.shape
			x = logits.view(B * T, C)
			y = targets.view(B * T)
			loss = F.cross_entropy(x, y)

		return logits, loss




# build the model
model = Gpt(batch_size, n_context, vocab_size, n_embd, n_layers, n_heads)
model = model.to(device)

# print the number of parameters in the model
print(sum(p.numel() for p in model.parameters())/1e6, 'M parameters')

# create a PyTorch optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

# optim loop
for epoch in range(max_iters):

	# every once in a while evaluate the loss on train and val sets
	if epoch % eval_interval == 0 or epoch == max_iters - 1:
		print(f"step {epoch}")

	# sample a batch of data
	xb, yb = get_batch('train')

	# evaluate the loss
	logits, loss = model(xb, yb)

	optimizer.zero_grad(set_to_none=True)
	loss.backward()
	optimizer.step()
	xm.mark_step()


</code></pre>