OOM with bfloat16, no speed-up about openfold HOT 14 CLOSED

It's not still missing in DeepSpeed. See current official documentation: https://www.deepspeed.ai/docs/config-json/#bfloat16-options.

As a first step, are you running DeepSpeed with ZeRO stage 1 or 2? Have you updated the DeepSpeed config in addition to passing bf16 as a parameter to the training script?

Our testing was done with the following:

CUDA Driver 465.19.01
CUDA 11.3 Update 1 (11.3.1.005)
cuBLAS 11.5.1.109 (part of CUDA 11.3 U1)
CUDNN 8.2.1.32
NCCL 2.9.9

As for the log_sigmoid thing, I don't know why that call would succeed with DeepSpeed and fail without it. As a sanity check, I'd try just spelling out the sigmoid formula instead.

from openfold.

Stage 2. My DeepSpeed config:

{ "optimizer": { "type": "Adam", "params": { "lr": 0.0005, "eps": 1e-05 } }, "amp": { "enabled": false, "opt_level": "O2" }, "bfloat16": { "enabled": true }, "zero_optimization": { "stage": 2, "cpu_offload": true, "contiguous_gradients": true }, "activation_checkpointing": { "partition_activations": true, "cpu_checkpointing": false, "profile": false }, "gradient_clipping": 0.1 }

Driver Version: 460.27.04 CUDA Version: 11.2

Is there no automatic casting with DeepSpeed? Features are still in torch.float32...

logsigmoid error without DeepSpeed disappears if I replace it with log(sigmoid(.)). In this version all losses go to NaN immediately and never recover. Same OOM. Features still in float32.

from openfold.

Just to clarify: you're also passing --precision bf16 to the training script, right?

from openfold.

Yes

from openfold.

In pure pytorch (no deepspeed, no lightning) I get the same logsigmoid error and all losses go to NaN immediately. There is no autocasting, features stay in torch.float32. Maybe because of the dict argument? I have to cast all floating point features manually.

What's interesting is, if I run deepspeed (with bfloat16 enabled) and manually cast the model to bfloat16 it complains: ValueError: fp32 is enabled but the following parameters have dtype that is not fp32: module.model.input_embedder.linear_tf_z_i.weight [..]

If I only use pytorch-lightning it actually shows in the beginning that it is in bfloat16 mixed precision mode (no message like this with deepspeed), but still no autocasting and NaN losses.

from openfold.

Hm. Interesting. I won't be personally up on A100s for a few more days, so there's not much I can do at the moment to investigate. I'll loop back when I do.

What happens if you manually cast both the model and the input features when you're using DeepSpeed? Do you think casting the model interferes with DeepSpeed's master FP32 copy of the weights? Maybe this would make a good DeepSpeed issue too.

from openfold.

"ValueError: fp32 is enabled but the following parameters have dtype that is not fp32: module.model.input_embedder.linear_tf_z_i.weight [..]"

This happens when I cast the model to bfloat32. The error is triggered in the setup phase in pytorch-lightning. When I only convert the data, the error occurs in the network, asking for a float instead.

I'm not sure what's happening exactly. I have the feeling that the bfloat16 flag in the configuration file is just ignored for some reason. Or it's some miscommunication between deepspeed and pytorch-lightning?

from openfold.

FYI: Earlier today I removed a line from the training script that silently changed the value of "CUDA_VISIBLE_DEVICES". This is a long shot, but could it be that this was moving computation onto GPUs that aren't compatible with bfloat16?

I'm not sure what's happening exactly. I have the feeling that the bfloat16 flag in the configuration file is just ignored for some reason. Or it's some miscommunication between deepspeed and pytorch-lightning?

Both of these sound plausible to me.

from openfold.

FYI: Earlier today I removed a line from the training script that silently changed the value of "CUDA_VISIBLE_DEVICES". This is a long shot, but could it be that this was moving computation onto GPUs that aren't compatible with bfloat16?

No, the nodes are in different queues.

from openfold.

I'm up on an A100 now, and I can replicate your NaN issue. I'll circle back if I can figure out how to fix it.

from openfold.

I'm up on an A100 now, and I can replicate your NaN issue. I'll circle back if I can figure out how to fix it.

This is when only using pytorch-lightning and changing the loss calculation? Does the bfloat16 flag in DeepSpeed have any effect for you? I piggybacked on the DeepSpeed issue above in the hope we get some response from the DeepSpeed people.

Thanks for following up!

from openfold.

It's just PyTorch Lightning, but I haven't changed the loss calculation. What do you mean by that?

from openfold.

It's just PyTorch Lightning, but I haven't changed the loss calculation. What do you mean by that?

Replacing logsigmoid with log(sigmoid(.)).

from openfold.

This should have been resolved by this week's commits. Closing this for now.

from openfold.