i am looking for any python code examples, either of following.
- model training library which I can specifcy abcType and computeType or scaleType all defined in cuda framework header files.
- or simply gemm3 example which also can specify these 3 entities.
So far I am able to do simple example for gemm3 but not sure how to specify abcType, computeType or scaleType differently in the example.
i cam up following snippet of code but not sure where or how to specify those:
I put this post in NLP, reason being usually in NLP tasks, speciflally inferencing these types are specified for lesser accuracy types to speed up inference.
import torch
import numpy as np
cuda = torch.device(‘cuda’)
m = 1024
n = 1024
k = 11024
t1 = torch.randn([m,n], dtype=torch.bfloat16, device=cuda)
t2 = torch.randn([n,k], dtype=torch.bfloat16, device=cuda)
print("t1, t2: ", t1.shape, t2.shape)
t3 = torch.mm(t1, t2)
#print(np.array(t3).shape)
print(t3.shape)