Incomplete transfer of module to device (only with 0.101.3)

[dje-dev]

Upon upgrading to 0.101.3 a previously working network now throws exception upon calling forward. The module is constructed on CPU with multiple layers in constructor and as normal, RegisterComponents is called. Module is subsequently moved to cuda. Most layers successfully are transferred, but at least one (LayerNorm) remains on cpu. This causes exception below.

Downgrading back to 0.101.2 immediately fixed the problem.

Unhandled exception. System.Runtime.InteropServices.ExternalException (0x80004005): Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument weight in method wrapper_CUDA__native_layer_norm)

[NiklasGustafsson]

@shaltielshmid -- I think this one will require your attention.

[shaltielshmid]

Can you please send sample code to replicate this behavior?

[shaltielshmid]

Okay, after some digging, I found the root of the problem to be from @GilesBathgate change here: #1152

In TorchSharp 0.101.3 the LayerNorm class was changed to act like torch.nn.functional.
In the constructor, the LayerNorm module calls the base constructor with a "IntPtr.Zero" handle.
When module.to() gets called, then it goes through all the registered modules and calls .to() on them as well. Because LayerNorm's handle is zero, so LayerNorm.handle.IsInvalid always returns true and therefore it is not registered as a module.

You can see the problematic line here.

@NiklasGustafsson Is there any reason to register the LayerNorm module with an invalid handle? Ideally we should just pass a string with the name of the module (either default it to "LayerNorm" or require it in the constructor).

[GilesBathgate]

@shaltielshmid if base(IntPtr.Zero, IntPtr.Zero) is problematic then it can easily be replaced with a named module base constructor. I chose to use IntPtr.Zero as this seems to be the way its done elsewhere, when there is no module handle required for calls to native functions. How does module.to work with user defined modules?

[GilesBathgate]

@NiklasGustafsson Perhaps I am meant to call base(IntPtr.Zero, IntPtr.Zero ownsHandle: false) ?

[shaltielshmid]

@GilesBathgate I'm not familiar with other modules that call base(IntPtr.Zero, IntPtr.Zero), can you send an example?

The problem with IntPtr.Zero is that the module is considered invalid and then it is not registered as a submodule of the main module. Aside from the module.to() issue this also leads to other problems, such as the LayerNorm.weight and LayerNorm.bias parameters not being returned in the module.parameters() function of the parent module. This means that the weights aren't passed to the optimizer and they aren't being tuned.

[GilesBathgate]

@shaltielshmid Literally every named user module, unless I am misunderstanding: https://github.com/dotnet/TorchSharp/blob/main/src/TorchSharp/NN/Module.cs#L965

[shaltielshmid]

@GilesBathgate In that case it's correct, since the constructor assigns the handle and boxedModule separately.

[NiklasGustafsson]

@shaltielshmid -- is the check for a valid handle on Module new? Every module that isn't implemented in native code will have an invalid native module handle.

[shaltielshmid]

@NiklasGustafsson No, it's always been there (well, at least since I started using TorchSharp) - located here.

But every module is given a valid handle of a custom module in the constructor here.

[GilesBathgate]

@shaltielshmid I see, yes, if I call base(IntPtr.Zero, IntPtr.Zero) I should also do handle = new HType(res, true); That was not what I expected. Fix would be like you say to call the named constructor., or add a new constructor that allows no name but inits handle accordingly.

[shaltielshmid]

@GilesBathgate Agreed.
I would vote for the named constructor option, since the Module is a custom module and probably best not to duplicate the code which creates the handle (and the other properties like _forwardNative).

[GilesBathgate]

I do wonder if Unnamed modules are a thing, and whether giving LayerNorm a name is compatible with PyTorch.

[NiklasGustafsson]

The names really don't have any functional significance.

[shaltielshmid]

To quote the documentation:

<param name="name">The name of the module. Useful for debugging purposes, mostly.</param>