RuntimeError: CUDA driver error: an illegal memory access was encountered about fms-fsdp HOT 9 CLOSED

It turns out that fsdp requires a slightly different approach when initializing after sharding. I'll open a PR to fix this on Monday

from fms-fsdp.

@daviswer technically, since you named it reset_parameters(), we can make param_init_fn=None as under the hood it will call reset_parameters() if no specific param_init_fn is passed. However, I vaguely remembered this wasn't working as expected last time (well, last time was end of last year, so maybe worth revisiting).

Can you prepare a small validation code snippet (to be called after FSDP call) to validate if the model is init as expected?

e.g. it should pass with current code, it should not pass with current code but removing model.reset_parameters(), and it should pass if we do a good param_init_fn.

from fms-fsdp.

@daviswer great. I will start working on this.

from fms-fsdp.

@daviswer did we revisit this thread? just making sure if this issue is still there.

from fms-fsdp.

I think #18 changed the calculus for how we were planning to handle this, and after that it never got revisited. Not sure if the issue is still relevant.

from fms-fsdp.

@daviswer yes. that pr helps this issue by moving reset_parameters() before FSDP call. But ultimately we would want the call to be made after the FSDP call, as we can save a full cpu model initialized before FSDP call, which can be 20 mins for large models like 70b.

So this will have to be fixed eventually.

from fms-fsdp.

So the way we'd want to do this is to add the (various) reset_parameters() to this portion of the FSDP call in main_training.py.

        param_init_fn=lambda module: (
            module.to_empty(device=torch.device("cuda"), recurse=False)
            if cfg.low_cpu_fsdp
            else None
        ),

But we need to make sure it keeps playing nicely with the the low_cpu_fsdp flag. Since you know that portion and I know the model init portion, we should probably coordinate @lchu-ibm

from fms-fsdp.

ok I opened a branch of fms main: fsdp_init_check, which adds a check_weights() function to Llama. This should error out for any improper init and return silently if successful

from fms-fsdp.

closing this one in favor of the new issue: #64

from fms-fsdp.