search

iree-org / iree-jax

Apache License 2.0
48 stars 19 forks source link

Make config.py confgurable from the command-line #68

Closed wangkuiyi closed 1 year ago

wangkuiyi commented 1 year ago

The goal of this pull request is to (1) preserve the present CI as is, and (2) allow users to export an MHLO for on-device inference and fine-tuning, which requires batch_size == 1.

I confirmed that I can execute the following commands to generate the MHLO file and build it into Metal GPU kernels.

python models/gpt2/export.py --batch_size=1 --no_compile

iree-compile \
  --iree-input-type=mhlo \
  --iree-hal-target-backends=metal \
  --iree-metal-compile-to-metallib=false \
  /tmp/gpt2.mlir \
  -o /tmp/gpt2-metal.vmfb

The command iree-compile was built from the feature branch https://github.com/antiagainst/iree/tree/apple-metal-hal.