Basically you'd familiarize yourself with transformers by implementing different variants of them, and changing them around according to your own ideas on different toy datasets.

Then you'd figure out a set of toy tasks that you like and think are important.

In this particular case you take something like NanoGPT, go to model.py, go to class GPT, go to __init__, modify the self.transformer ModuleDict by changing nn.Embedding to a ModuleList of nn.Embedding, then you change the for loop at line 180 to loop over a range, modify forward by adding x = x + self.transformer.wte[i], something like that I think.

I haven't tried yet though (I've got a terrible cold, so I am on social media instead of doing anything sensible).

Also, this particular thing didn't work on my toy problems. It might still be good though.