Thanks!

[danielhanchen]

Hey! Much apologies on the delay - fantastic work - sadly I've been inundated with stuff, so I can't review your submission in full - in general it's an extremely strong submission, so appreciate the effort and great work again!

I'll get back to you in a few days if that works! Your submission definitely was top notch!

[RameshBabuAsh]

Thanks for the update, even though you haven’t had a chance to review everything yet, I really appreciate the feedback. I’m excited about the opportunity and look forward to hearing from you soon :)

[danielhanchen]

Apologies on the massive delay! Thank you for your patience and great work again!

1. Task A - unfortunately you have mismatched elements, so this means the kernel is wrong.

AssertionError: Tensor-likes are not close!

Mismatched elements: 9 / 3978240 (0.0%)
Greatest absolute difference: 0.015625 at index (1, 353, 453) (up to 0.01 allowed)
Greatest relative difference: 0.921875 at index (1, 452, 768) (up to 0.01 allowed)


During handling of the above exception, another exception occurred:

2. Task B - FSDP2: This is not FSDP2 but FSDP!
3. Task C - Torch Compile: Great work, but unfortunately the training losses are not correct (for both uncompiled and compiled). Also I think there's only 10 compiled debug messages - this seems odd - there should be say 20 at least!
4. Task E - GRPO is incorrect! Using nn.Linear = MemoryEfficientLinear is definitely unfortunately wrong - this patches over ALL linear layers, which makes training slower, and is not advised. Also in your Llama 1B check:

# Print each layer's name along with its trainable status.
for name, param in model.named_parameters():
    if "model.layers.15" not in name:
        param.requires_grad = False

why is the gradient frozen except for layer 15?

[danielhanchen]

Actually on that note - on FSDP on Kaggle - I'm seeing very similar solutions to your one like here: https://www.kaggle.com/code/nuezzz/challenge-b - are these FSDP solutions from a source that I should know about? Or this is your original work and someone else unfortunately did not credit your work? Thanks!

[RameshBabuAsh]

Thanks for the feedback @danielhanchen !

1. When the elements are closer to zero the kernel goes wrong, which results in the greatest absolute difference of about 0.9. I’ll need to work on my Triton knowledge here, as I’m still relatively new to it.

2. I mistakenly assumed that accelerate supported FSDP2 natively - it turns out it does not. Just to clarify, this is entirely my work rather than derived from any external source and that is the reason for the blunder as well!

3. Regarding the Torch Compile debug messages: the fewer messages, the better the graphs, right? I’ve observed that the loss is unstable in the provided problem notebook as well, similar to what I encountered in my submission. I’m still investigating the root cause of this instability.

4. For the Llama 1B check, I replaced all the linear layers with MemoryEfficientLayer to ensure that the losses match, because the unquantized version couldn’t fit in memory for complete training. This is why I froze all layers except for layer 15. I’ve trained using both nn.Linear and MemoryEfficientLayer separately and demonstrated that all losses match accordingly.