Efficient Sampling

[eleutherai]

Efficient Sampling

[Mistobaan]

@shawwn is your pr #28 related to this?

[StellaAthena]

Per Stack Overflow, apparently you cannot have multiple assignees to a private repo. No idea why.

@Mistobaan @lucidrains should assure you both get notifications.

[sdtblck]

reopening this because there are still some problems

1. at least with the models I've tested, 'fast sampling' seems to be of significantly worse quality than slow sampling. (see sample screenshots below)

Slow sampling:

Fast sampling:

It looks ok on the surface but is notably more nonsensical and more often than not devolves into repetition:

2. fast sampling does not work with larger TPU pods - I will post the stacktrace for this shortly (whenever we have a spare pod 😆 )

I'm happy to look into this if you're busy @lucidrains . I suspect 1) is just an off by one error somewhere with the attention mask, maybe. 2) might be a little trickier but I'd be happy to look into this one, since TPUs might be the one area where I can claim to have more experience, haha. I think it may be a shape error relating to the mtf splitting.

[sdtblck]

After a bit more digging, it seems like the sample quality is only worse with MOE models. I'm going to keep investigating as to why.

[sdtblck]

Ok, for now we have set fast sampling as the default, but revert to slow sampling for moe models.

Error no. 2 wasn't in fact an error with larger pods, or models, but was an error with fast sampling when recompute_grad was set to True. This is fixed here: 8415e56

so what's left now is to 1) figure out what's wrong with moe and 2) try out step 4 here: #43

[sdtblck]

closing this issue to open a more relevant one w/r/t moe