is mx.load lazy eval?

[mzbac]

From my understanding, the mx.load function is doing lazy evaluation, otherwise we wouldn't be able to quant large LLMs. For example, I am having no issues quantise the llama 3.1 model with 405B which requires over 700GB of memory if it loads the entire weight into memory without lazy eval that won't work. Does this mean that in the code snippet provided below:

weights = {}
for wf in weight_files:
    weights.update(mx.load(wf))

The weight dict is not fully loaded until we call weight.items. Thank you in advance.

[angeloskath]

Yes exactly. Until you actually run mx.eval(weights) no array has been read from the files (metadata have been read in order to know how many arrays, shapes, types etc).

There are some caveats with lazy loading. If a python file-like object is provided, then load evaluates the arrays immediately. The reason is we cannot be sure that the object will be valid after calling load. For instance the following code will not be lazy

with open(weight_file) as f:
    weights = mx.load(f)
# we couldn't wait cause the file here is closed

[mzbac]

Yes exactly. Until you actually run mx.eval(weights) no array has been read from the files (metadata have been read in order to know how many arrays, shapes, types etc).

There are some caveats with lazy loading. If a python file-like object is provided, then load evaluates the arrays immediately. The reason is we cannot be sure that the object will be valid after calling load. For instance the following code will not be lazy

with open(weight_file) as f:
    weights = mx.load(f)
# we couldn't wait cause the file here is closed

Thanks @angeloskath, really appreciate the clarification.