Issue when upgrading from MLX version 0.4.0 to 0.6.0

[manjosk8]

Hi,

I'm encountering internal MLX errors after upgrading to version 0.6.0. The same code functions properly with version 0.4.0. I attempted to address this by calling mx.eval(model.parameters()) before method for m in model.modules(), but unfortunately, it didn't resolve the issue. Are there any new global flags introduced in version 0.6.0 that I should be aware of before running the project? For instance, something akin to mx.set_default_device(set_device) that could help maintain the same behavior as in version 0.4.0.

In version 0.6.0, mx.eval(model.parameters()) by itself returns the same error:

File "/myproject/utils/torch_utils.py", line 395, in initialize_weights_mlx for m in model.modules(): File ".../venv/lib/python3.10/site-packages/mlx/nn/layers/base.py", line 439, in modules self.apply_to_modules(lambda k, m: modulelist.append(m)) File ".../venv/lib/python3.10/site-packages/mlx/nn/layers/base.py", line 429, in apply_to_modules tree_flatten(mod.children(), prefix=prefix, is_leaf=self.is_module) File ".../venv/lib/python3.10/site-packages/mlx/nn/layers/base.py", line 305, in children return self.filter_and_map( File ".../venv/lib/python3.10/site-packages/mlx/nn/layers/base.py", line 287, in filter_and_map return { File ".../venv/lib/python3.10/site-packages/mlx/nn/layers/base.py", line 288, in <dictcomp> k: _unwrap(self, k, v, filter_fn, map_fn, is_leaf_fn) File ".../venv/lib/python3.10/site-packages/mlx/nn/layers/base.py", line 23, in _unwrap for k, v in v.items(): UnboundLocalError: local variable 'v' referenced before assignment

Regards,
Marijan

[awni]

Nothing should have broken from 0.4 to 0.6, this is likely a bug. Could you provide a way for us to reproduce the issue?

[manjosk8]

Hi Awni,

It's challenging for me to share the proprietary code, but I can offer you a hint. I've constructed a network using nn.Sequential, and it worked without any problems for versions up to 0.4.0. However, after upgrading to 0.5.0 or 0.6.0, the issue mentioned above arises.

Initially, I suspected that the error might be due to forgetting to call mx.eval(module.parameters()) after netowork is constructed, but that doesn't seem to be the case. Even after ensuring that I call mx.eval, the issue persists for versions beyond 0.4.0., library throws exception while executing mx.eval(module.parameters()) so that's something internal.

There appears to be a significant change in the library's internals between versions 0.4.0 and later, causing this issue.

Marijan

[awni]

There is no need for you to share you proprietary code, just a simple test case that reproduces the issue would be good (in fact even much better!). I really have no idea what the problem is just by looking at the stack trace and it will be quite difficult to debug. nn.Sequential seems to work fine in a basic test:

import mlx.core as mx
import mlx.nn as nn

model = nn.Sequential(nn.Linear(2, 4), nn.ReLU(), nn.Linear(4, 1))
mx.eval(model.parameters())
output = model(mx.zeros((8, 2)))
print(output)

[awni]

Actually scratch that, I found the bug :)

[awni]

Fix here: #819

[manjosk8]

Hi Awni,

thank you for fixing the previous bug I can confirm that the feature works, but I've found new issues with MLX 0.7.0 and they happen when I try to save model using model.save_weights('some_array.npz') I am getting the following internal errors and they don't occur on MLX 0.4.0, as before same code new version of MLX lib 0.4.0 works without any issues and 0.7.0 throws:

When saving model class inherited from class that inherits nn.Module:

File ".../lib/python3.10/site-packages/mlx/nn/layers/base.py", line 231, in save_weights
mx.savez(file, **params_dict)
ValueError: [save] cannot serialize an empty array

When saving nn.Sequential model:
File ".../lib/python3.10/site-packages/mlx/nn/layers/base.py", line 231, in save_weights
mx.savez(file, **params_dict)
ValueError: [save] can only serialize row or col contiguous arrays

EDITED:
Found issue with an empty array it happens due to the property locking of MLX when we have instance of nn.Module class and we initialize variable in costructor to e.g.: self.a = mx.array([]), even if during the runtime we change self.a to be valid array in one method e.g. self.a = mx.array([1,2]) MLX never updates the value it always remains mx.array([]). MLX 0.4.0 does not behave like that

For second issue it remains unresolved and we'll try migrating to numpy.savez for now

[awni]

For the first issue:

even if during the runtime we change self.a to be valid array in one method e.g. self.a = mx.array([1,2]) MLX never updates the value it always remains mx.array([])

The following does exactly what you expect (updates a to the array). I suspect something is up with how you are setting the array. If you are able to share a slightly more detailed example that reproduces the problem I can help debug.

import mlx.core as mx
import mlx.nn as nn

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.a = mx.array([])

    def __call__(self):
        pass

    def set(self):
        self.a = mx.array([1,2])

model = MyModule()
model.set()

print(model.parameters())

[awni]

ValueError: [save] can only serialize row or col contiguous arrays

You shouldn't be getting that. I will see if I can reproduce it.

[awni]

Are you able to provide a little code to reproduce the serialization contiguity error?