error:tImplementedError: Could not run 'aten::_amp_foreach_non_finite_check_and_unscale_' with arguments from the 'CPU' backend.

[daixiangzi]

when I use:
from torch.cuda.amp import GradScaler
amp = GradScaler(init_scale=512, growth_interval=100)
for img, label in train_iter:
with torch.cuda.amp.autocast(True):
label = backbone(img)
amp.scale(loss).backward()
amp.unscale_(opt)

[lucidrains]

@daixiangzi that looks like an unrelated error that is specific to cpu with the amp.scale? could you retry on cuda?

[daixiangzi]

@daixiangzi that looks like an unrelated error that is specific to cpu with the amp.scale? could you retry on cuda?

When I use a single machine with multiple gpu, I can run your code, but when I use multiple machines with multiple gpus,error will appear.

[daixiangzi]

@daixiangzi that looks like an unrelated error that is specific to cpu with the amp.scale? could you retry on cuda?
I have always used it on cuda

[lucidrains]

are you using fsdp?

[daixiangzi]

fsdp

no ,I use ddp

[lucidrains]

ah, hard for me to debug. i only have a single machine

[lucidrains]

@daixiangzi could you try disabling mixed precision just for the MoE block?

with torch.cuda.amp.autocast(enabled = False):
    # ... your moe forward

[lucidrains]

finally seeing the limits of pytorch

in jax, properly working moe is just a few lines of code

[daixiangzi]

@daixiangzi could you try disabling mixed precision just for the MoE block?

with torch.cuda.amp.autocast(enabled = False):
    # ... your moe forward

ok

[daixiangzi]

@daixiangzi could you try disabling mixed precision just for the MoE block?

with torch.cuda.amp.autocast(enabled = False):
    # ... your moe forward

I try it just now, but not work

[daixiangzi]

define moe layer:

self.softmoe = SoftMoE(dim = 512,num_experts =64,is_distributed=True,offload_unused_experts_to_cpu=True)

forward:

with torch.cuda.amp.autocast(enabled =False):
                x = x + self.softmoe(self.ln_3(x).permute(1,0,2)).permute(1,0,2)

[daixiangzi]

NotImplementedError: Could not run 'aten::amp_foreach_non_finite_check_and_unscale' with arguments from the 'CPU' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build proces

[daixiangzi]

in fact I train dp version on single machine with multiple gpu, but result is still bad, I have a little doubt about

soft-moe-pytorch/soft_moe_pytorch/soft_moe.py

Line 315 in 4dc07b0

x = self.norm(x)

I see origin rep:https://github.com/google-research/vmoe/blob/dfb9ee01ce6dfc5a8228b406e768ee325dd18fcd/vmoe/nn/vit_moe.py#L109C27-L109C45
your code:

x = self.norm(x)
logits = einsum('b n d, e s d -> b n e s', x, slot_embeds)
...
slots = einsum('b n d, b n e s -> b e s d', x, dispatch_weights)

origin rep code seem is:

norm_x = self.norm(x)
logits = einsum('b n d, e s d -> b n e s', norm_x, slot_embeds)
....
slots = einsum('b n d, b n e s -> b e s d', x, dispatch_weights)

[lucidrains]

@daixiangzi hmm, have you tried it their way? conventionally we always work with pre-normalized values, so that would be weird they dispatch based on unnormalized input into the block

[daixiangzi]

@daixiangzi hmm, have you tried it their way? conventionally we always work with pre-normalized values, so that would be weird they dispatch based on unnormalized input into the block

no ,I will try it.but I always have some doubts about the effectiveness of softmoe, mainly because I have been adjusting it for a long time

[daixiangzi]

In terms of visual tasks, there are currently many open source versions of MOE, but in fact, there is not a sufficient experiment to demonstrate the effectiveness of this approach

[lucidrains]

@daixiangzi yea it happens

thanks for sharing your results

just to make sure we don't miss anything, would you like to also do a quick run where i substituted the rmsnorm with layernorm? (set https://github.com/lucidrains/soft-moe-pytorch/blob/main/soft_moe_pytorch/soft_moe.py#L287 to True)

[daixiangzi]

@daixiangzi yea it happens

thanks for sharing your results

just to make sure we don't miss anything, would you like to also do a quick run where i substituted the rmsnorm with layernorm? (set https://github.com/lucidrains/soft-moe-pytorch/blob/main/soft_moe_pytorch/soft_moe.py#L287 to True)

Actually, before this, I tried my own version where the norm part used LN, but the effect was still not as good as the baseline

[lucidrains]

@daixiangzi ah ok, good to know

i'll take your experience as a datapoint

[daixiangzi]

myself softmoe version:

        #x: [b,50,768]
        #phi [768, 128, slots]
        phi = self.phi
        phi = self.scale*self.normalize(phi,axis=0)
        #logits:[b, 50, 128, 1]
        
        logits = einsum(self.normalize(x,axis=-1), phi, "b m d, d n p -> b m n p")
        #add noise
        if self.training:
            normal = torch.randn(logits.shape, dtype=logits.dtype,device=logits.device)
            logits = logits +(1.0 * normal)
        #dispatch_weights:[b, token_num, 128, 1])
        dispatch_weights = logits.softmax(dim=1)  # denoted 'D' in the paper
        # NOTE: The 'torch.softmax' function does not support multiple values for the
        # 'dim' argument (unlike jax), so we are forced to flatten the last two dimensions.
        # Then, we rearrange the Tensor into its original shape.
        #对experts softmax:combine_weights:[b, token_num, 128, 1]
        combine_weights = rearrange(logits.flatten(start_dim=2).softmax(dim=-1),"b m (n p) -> b m n p",n=self.num_experts)

        # NOTE: To save memory, I don't rename the intermediate tensors Y, Ys, Xs.
        # Instead, I just overwrite the 'x' variable.  The names from the paper are
        # included in a comment for each line below.
        
        x = einsum(x, dispatch_weights, "b m d, b m n p -> b n p d")  # Xs
        #x:[b,128, 1, out_d]
        x = self.experts(x)  # Ys
        #combine_weights:[b, 50, 128, 1]
        x = einsum(x, combine_weights, "b n p d, b m n p -> b m d")  # Y
        #[b, 128, 1, 384]
        return x

[daixiangzi]

in softmoe paper:
self.normalize should l2_norm, for example:

 def normalize(self,x:torch.Tensor,axis:int):
        return x*torch.rsqrt(torch.square(x).sum(dim=axis, keepdim=True) + 1e-6)

[lucidrains]

in softmoe paper: self.normalize should l2_norm, for example:

 def normalize(self,x:torch.Tensor,axis:int):
        return x*torch.rsqrt(torch.square(x).sum(dim=axis, keepdim=True) + 1e-6)

ohh i did it correct then the first time

[daixiangzi]

in softmoe paper: self.normalize should l2_norm, for example:

 def normalize(self,x:torch.Tensor,axis:int):
        return x*torch.rsqrt(torch.square(x).sum(dim=axis, keepdim=True) + 1e-6)

ohh i did it correct then the first time

ok