Compilation: [eval] Attempting to eval an array without a primitive.

[danilopeixoto]

Can a training step with Mixtral model be compiled?

state = [model.state, optimizer.state]

@partial(mx.compile, inputs=state, outputs=state)
def train_step(batch):
    loss_value_and_grad = nn.value_and_grad(model, default_loss)

    (lvalue, toks), grad = loss_value_and_grad(model, *batch)
    optimizer.update(model, grad)

    return lvalue.item(), toks.item()

When attempting to compile, the training step yields the following error:

File "mixtral.py", line 172, in __call__
    mx.eval(inds)
ValueError: [eval] Attempting to eval an array without a primitive.

[awni]

What code are you using there? It should be compilable, but you have to be careful to make sure all the implicit state is captured. Usually you see that error message when you forget to include state in the inputs and/or outputs.

That said I don't think you will see much gain (yet) from compiling it since most of the work should be in the matrix multiplications in the MLPs and attention and that is not affected from compile.

[danilopeixoto]

What code are you using there?

The model or mixtral.py is ml-explore/mlx-examples/blob/main/llms/mlx_lm/models/mixtral.py file.

The default_loss is:

https://github.com/ml-explore/mlx-examples/blob/f71e965d57c01a2b4ac3d43d72f75482dc80fe86/llms/mlx_lm/tuner/trainer.py#L43

[awni]

Ah sorry I should have realized earlier, you cannot compile the MOE models right now as they do an implicit graph eval to determine the expert to route to. That needs a workaround which we have not implemented yet.

[stefanvarunix]

Hi,
I am not sure of I get the problem right,
running
python -m mlx_lm.lora --model /pathtoMixtral --train --data /pathtodata --iters 600
gives the error
ValueError: [eval] Attempting to eval an array without a primitive.

from /mlx-examples/lora, running
python lora.py --model /pathtoMixtral --train --data /pathtodata --iters 600
starts the training without that error.

The lora.py from /mlx-examples/lora is older than the lora.py from mlx_lm. The new version should be the version of choice. or?
Anything I can do?

mlx Version = 0.7.0

[mzbac]

The latest version of mlx-lm introduced the compile in the training step, which does not work with Mixtral model. You can try removing this line from the source code or downgrading mlx-lm.

[awni]

Hmm, I thought I disabled compile if it's an MOE model.. I must have missed something there.