def test_embedding_bag_basic(self):
"""Test of aten::embedding_bag node on glow"""
check_skip(self)
class TestModule(torch.nn.Module):
def forward(self, input, offsets, per_sample_weights):
weight = torch.FloatTensor([[1, 2.3, 3], [4, 5.1, 6.3]])
embedding_sum = torch.nn.EmbeddingBag.from_pretrained(
weight, mode="sum")
# in jit mode we need to discard the end offset
a = embedding_sum(input, offsets[:-1])
b = embedding_sum(input, offsets[:-1], per_sample_weights)
return a, b
input = torch.LongTensor([1, 0, 0, 1, 1])
offsets = torch.LongTensor([0, 1, 5]) # final item is endOffset
per_sample_weights = torch.FloatTensor([1, 2, 3, 4, 5])
utils.compare_tracing_methods(
TestModule(),
input,
offsets,
per_sample_weights,
fusible_ops={"aten::embedding_bag"},
skip_to_glow=
True, # to_glow doesn't support include_last_offset=False
)
def test_embedding_3d_indices(self):
"""Test of aten::embedding node on glow"""
check_skip(self)
class TestModule(torch.nn.Module):
def forward(self, indices):
weight = torch.FloatTensor([[1, 2.3, 3], [4, 5.1, 6.3], [3.3, -4, 5.7]])
embedding = torch.nn.Embedding.from_pretrained(weight)
a = embedding(indices)
return a
indices = torch.LongTensor([[[1, 2], [0, 1]]])
utils.compare_tracing_methods(
TestModule(),
indices,
fusible_ops={"aten::embedding"},
skip_to_glow=True, # to_glow doesn't support include_last_offset=False
)
def test_embedding_bag_basic(self):
"""Test of aten::embedding_bag node on glow"""
check_skip(self)
def embedding_bag_basic(input, offsets, per_sample_weights):
weight = torch.FloatTensor([[1, 2.3, 3], [4, 5.1, 6.3]])
embedding_sum = torch.nn.EmbeddingBag.from_pretrained(weight, mode="sum")
# in jit mode we need to discard the end offset
a = embedding_sum(input, offsets[:-1])
b = embedding_sum(input, offsets[:-1], per_sample_weights)
return a, b
input = torch.LongTensor([1, 0, 0, 1, 1])
offsets = torch.LongTensor([0, 1, 5]) # final item is endOffset
per_sample_weights = torch.FloatTensor([1, 2, 3, 4, 5])
jitVsGlow(
embedding_bag_basic,
input,
offsets,
per_sample_weights,
expected_fused_ops={"aten::embedding_bag"},
)
def test_embedding_bag_rowwise_offsets(self, name, num_lengths, is4bit,
is_weighted, use_fp16):
"""Test of quantized::embedding_bag_byte_rowwise_offsets and
quantized::embedding_bag_4bit_rowwise_offsets node on glow"""
check_skip(self)
class TestModule(torch.nn.Module):
def __init__(self,
q_weights,
is4bit=False,
per_sample_weights=None):
super().__init__()
self.q_weights = q_weights
self.per_sample_weights = per_sample_weights
if is4bit:
self.op = torch.ops.quantized.embedding_bag_4bit_rowwise_offsets
else:
self.op = torch.ops.quantized.embedding_bag_byte_rowwise_offsets
def forward(self, indices, offsets):
return self.op(
self.q_weights,
indices,
offsets,
mode=0,
per_sample_weights=self.per_sample_weights,
include_last_offset=True,
)
# generate random weights and indices
num_embeddings = 16
embedding_dim = 4
weights = torch.from_numpy((np.random.random_sample(
(num_embeddings, embedding_dim)) + 1).astype(np.float32))
q_weights = (torch.ops.quantized.embedding_bag_4bit_prepack(weights)
if is4bit else
torch.ops.quantized.embedding_bag_byte_prepack(weights))
np_lengths = (np.zeros(shape=[10]).astype(np.int32)
if num_lengths == 0 else np.random.randint(
0, num_lengths, size=10).astype(np.int32))
num_lengths = np.sum(np_lengths)
lengths = torch.from_numpy(np_lengths)
indices = torch.from_numpy(
np.random.randint(low=0,
high=num_embeddings,
size=num_lengths,
dtype=np.int64)).long()
offsets = torch.cat([torch.zeros([1]),
torch.cumsum(lengths, 0)]).long()
per_sample_weights = torch.from_numpy(
np.random.uniform(low=0.01, high=0.5,
size=[len(indices)]).astype(np.float32))
m = TestModule(q_weights, is4bit,
per_sample_weights if is_weighted else None)
utils.compare_tracing_methods(
m,
indices,
offsets,
fusible_ops={
"quantized::embedding_bag_4bit_rowwise_offsets"
if is4bit else "quantized::embedding_bag_byte_rowwise_offsets"
},
fp16=use_fp16,
# FP16 version is known to yeild different results, so our
# test here is mainly focusing on the flow rather than actual
# accuracy. There will be additional coverage on accuracy of
# the lowered modules
atol=0.02 if (is4bit or use_fp16) else 5e-4,
)
def test_glow_embedding_bag_rowwise_offsets(
self,
name,
num_lengths,
is4bit,
is_weighted,
use_fp16,
per_sample_weights_fp16,
):
"""Test of fb::glow_embedding_bag_byte_rowwise_offsets"""
check_skip(self)
class TestModule(torch.nn.Module):
def __init__(self, is4bit=False, per_sample_weights=None):
super().__init__()
self.per_sample_weights = per_sample_weights
if is4bit:
self.op = torch.ops.fb.glow_embedding_bag_4bit_rowwise_offsets
else:
self.op = torch.ops.fb.glow_embedding_bag_byte_rowwise_offsets
def forward(self, indices, offsets):
# Tensor input, Tensor? offsets, str weight_qualname, int num_embeddings, int embedding_dim, Tensor? per_sample_weights=None, bool include_last_offset=False
return self.op(
indices, # Tensor input
offsets, # Tensor? offsets
"weight_id", # str weight_qualname
16,
4,
self.per_sample_weights,
True,
)
num_embeddings = 16
np_lengths = (np.zeros(shape=[10]).astype(np.int32)
if num_lengths == 0 else np.random.randint(
0, num_lengths, size=10).astype(np.int32))
num_lengths = np.sum(np_lengths)
lengths = torch.from_numpy(np_lengths)
indices = torch.from_numpy(
np.random.randint(low=0,
high=num_embeddings,
size=num_lengths,
dtype=np.int64)).long()
offsets = torch.cat([torch.zeros([1]),
torch.cumsum(lengths, 0)]).long()
per_sample_weights_type = (np.float16 if per_sample_weights_fp16
and is4bit else np.float32)
per_sample_weights = torch.from_numpy(
np.random.uniform(low=0.01, high=0.5,
size=[len(indices)
]).astype(per_sample_weights_type))
m = TestModule(is4bit, per_sample_weights if is_weighted else None)
with self.assertRaisesRegex(
RuntimeError,
"Error Placholder: weight_id is marked as static but no deferredWeightLoader is provided.",
):
utils.test_lowering(
m,
indices,
offsets,
fusible_ops={
"fb::glow_embedding_bag_4bit_rowwise_offsets" if is4bit
else "fb::glow_embedding_bag_byte_rowwise_offsets"
},
fp16=use_fp16,
scripted=True,
)
def test_embedding_bag_rowwise_offsets(
self, name, num_lengths, is4bit, is_weighted, use_fp16
):
"""Test of quantized::embedding_bag_byte_rowwise_offsets and
quantized::embedding_bag_4bit_rowwise_offsets node on glow"""
check_skip(self)
class TestModule(torch.nn.Module):
def __init__(self, q_weights, is4bit=False, per_sample_weights=None):
super().__init__()
self.q_weights = q_weights
self.per_sample_weights = per_sample_weights
if is4bit:
self.op = torch.ops.quantized.embedding_bag_4bit_rowwise_offsets
else:
self.op = torch.ops.quantized.embedding_bag_byte_rowwise_offsets
def forward(self, indices, offsets):
return self.op(
self.q_weights,
indices,
offsets,
mode=0,
per_sample_weights=self.per_sample_weights,
include_last_offset=True,
)
# generate random weights and indices
num_embeddings = 16
embedding_dim = 4
weights = torch.from_numpy(
(np.random.random_sample((num_embeddings, embedding_dim)) + 1).astype(
np.float32
)
)
q_weights = (
torch.ops.quantized.embedding_bag_4bit_prepack(weights)
if is4bit
else torch.ops.quantized.embedding_bag_byte_prepack(weights)
)
np_lengths = (
np.zeros(shape=[10]).astype(np.int32)
if num_lengths == 0
else np.random.randint(0, num_lengths, size=10).astype(np.int32)
)
num_lengths = np.sum(np_lengths)
lengths = torch.from_numpy(np_lengths)
indices = torch.from_numpy(
np.random.randint(
low=0, high=num_embeddings, size=num_lengths, dtype=np.int64
)
).long()
offsets = torch.cat([torch.zeros([1]), torch.cumsum(lengths, 0)]).long()
per_sample_weights = torch.from_numpy(
np.random.uniform(low=0.01, high=0.5, size=[len(indices)]).astype(
np.float32
)
)
m = TestModule(q_weights, is4bit, per_sample_weights if is_weighted else None)
jitVsGlow(
m.forward,
indices,
offsets,
expected_fused_ops={
"quantized::embedding_bag_4bit_rowwise_offsets"
if is4bit
else "quantized::embedding_bag_byte_rowwise_offsets"
},
use_fp16=use_fp16,
# Known issue: 4bit version has bigger error
atol=0.005 if is4bit else 5e-4,
backend_name=get_backend_name(),
)
