Hi,
I went over the really good content on the MAX LLM book but I can’t seem to get the model to compile. To simplify the problem right now I’m just trying to get the MLP part of the model to compile but without any success.
Here is the code to reproduce the error
import max.functional as F
from max.graph import DeviceRef
from max.nn import Linear, Module
from max.tensor import Tensor, TensorType, defaults
class MLP(Module):
def __init__(self, emb: int, factor: int) -> None:
super().__init__()
self.up = Linear(in_dim=emb, out_dim=factor * emb, bias=True)
self.down = Linear(in_dim=factor * emb, out_dim=emb, bias=True)
def forward(self, x: Tensor) -> Tensor:
# x: [B,S,D]
x = self.up(x)
x = F.gelu(x, approximate="tanh")
x = self.down(x)
return x
def main() -> None:
dtype, device = defaults()
print(f"{dtype=} | {device=}")
model = MLP(emb=10, factor=2)
print(model)
x_type = TensorType(dtype, (5, 10), device=DeviceRef.from_device(device))
print("Compiling")
compiled_model = model.compile(x_type)
if __name__ == "__main__":
main()
I get the following error message
dtype=bfloat16 | device=Device(type=gpu,id=0)
MLP(
up=Linear(in_dim=Dim(10), out_dim=Dim(20)),
down=Linear(in_dim=Dim(20), out_dim=Dim(10))
)
Compiling
Traceback (most recent call last):
File "/scratch/ap6604/.pixi/envs/default/lib/python3.13/site-packages/max/engine/api.py", line 424, in load
_model = self._impl.compile_from_object(
model._module._CAPIPtr, # type: ignore
custom_extensions_final,
model.name,
)
ValueError: Failed to run the MOToMGP pass manager:
-:1:1: error: failed to lower module to LLVM IR for archive compilation, run LowerToLLVMPipeline failed
error: invalid rebind between two unequal types: !pop.array<2, scalar<si32>> to !pop.array<2, scalar<si64>>
error: failed to legalize operation 'kgen.rebind' that was explicitly marked illegal
-:1:1: error: The graph compiler could not elaborate the generated KGEN
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/home/ap6604/code_examples/modular/models/main.py", line 32, in <module>
main()
~~~~^^
File "/home/ap6604/code_examples/modular/models/main.py", line 28, in main
compiled_model = model.compile(x_type)
File "/scratch/ap6604/.pixi/envs/default/lib/python3.13/site-packages/max/nn/module.py", line 654, in compile
compiled = F.functional(session.load(graph, weights_registry=weights))
~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/scratch/ap6604/.pixi/envs/default/lib/python3.13/site-packages/max/engine/api.py", line 430, in load
raise RuntimeError(
...<5 lines>...
) from e
RuntimeError: Failed to compile the model. Please file an issue, all models should be correct by construction and this error should have been caught during construction.
For more detailed failure information run with the environment variable `MODULAR_MAX_DEBUG=True`.
I’m currently using:
- NVIDIA H200 GPU
- Mojo 0.26.2.0.dev2026021105 (1982f891)
- Ubuntu 24.04.3 LTS
It would be great If someone could try to compile the code above and share the pixi.toml!
