import mindspore

from mindspore import nn, ops

from mindspore import numpy as mnp



from model import Test



class ParallelConfig:

def __init_(self, data_parallel=1, model_parallel=1):

self.data_parallel = data_parallel

self.model_parallel = model_parallel

if __name__ ==  “__main__”:

mindspore.set_context(mode=mindspore . PYNATIVE_MODE)

mindspore.set_auto_parallel_context(device_num=2)

mindspore.set_auto_parallel_context(parallel_mode=”auto_parallel”)

mindspore.set_auto_parallel_context(search_mode=”sharding_propagation”)

batch_size = 2

dim = 8

n_heads = 4

input = mnp.ones((batch_size, 16, dim))

Config=ParallelConfig(data_parallel=1, model_parallel=1)

attn = Test(dim=dim, n_heads=n_heads, parallel_config=config)

output = attn(input)

print(output)





import mindspore

from mindspore import nn, ops

from mindspore import numpy as mnp



class Test(nn.Cell):

def _init_(self,

dim=512,

n_layers=8,

n_head=8,

vocab_size=-1, # defined Later by tokenizer

multiple_of=256, # make SWIGLU hidden Layer size muLtipLe of Large power of 2

norm_eps=1e-5,

max_batch_size=32,

max_seq_len=2048,

parallel_config=None,):

super()._init_()

self.n_local_heads = n_heads // 2

self.head_dim = dim // n_heads

self.w_q = nn.Dense( dim, n_heads * self.head_dim, has_bias = False)

self.w_q.shard(in_strategy((parallel_config.data_parallel, 1, 1),),

out_strategy(None,),

parameter_plan{”self.w_q.weight”:(parallel_config.model_parallel,1)},

device=”GPU”)



def construct(self, _x):

bsz, seqien, _ = _x.shape

x_q = self.w_q(_x)

return x_q