Created
February 21, 2023 23:44
-
-
Save wangkuiyi/9be63c184366d91f4f3ecc4c8ef08f66 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 14:03 $ python export.py | |
| I0221 15:44:29.689399 8453013824 xla_bridge.py:170] Remote TPU is not linked into jax; skipping remote TPU. | |
| I0221 15:44:29.689505 8453013824 xla_bridge.py:421] Unable to initialize backend 'tpu_driver': Could not initialize backend 'tpu_driver' | |
| I0221 15:44:29.689548 8453013824 xla_bridge.py:421] Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig' | |
| I0221 15:44:29.689573 8453013824 xla_bridge.py:421] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig' | |
| I0221 15:44:29.689785 8453013824 xla_bridge.py:421] Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client' | |
| I0221 15:44:29.689874 8453013824 xla_bridge.py:421] Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this. | |
| Traceback (most recent call last): | |
| File "/Users/y/w/iree-ios/iree-jax/models/gpt2/export.py", line 110, in <module> | |
| absl.app.run(main) | |
| File "/Users/y/miniforge3/envs/transformers/lib/python3.10/site-packages/absl/app.py", line 308, in run | |
| _run_main(main, args) | |
| File "/Users/y/miniforge3/envs/transformers/lib/python3.10/site-packages/absl/app.py", line 254, in _run_main | |
| sys.exit(main(argv)) | |
| File "/Users/y/w/iree-ios/iree-jax/models/gpt2/export.py", line 101, in main | |
| f.write(str(Program.get_mlir_module(module))) | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/program_api.py", line 377, in get_mlir_module | |
| info = Program.get_info(Program._get_instance(m)) | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/program_api.py", line 372, in _get_instance | |
| m = m() | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/program_api.py", line 459, in __new__ | |
| export_function() | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/program_api.py", line 454, in export_function | |
| info.export_module.def_func(invoke_with_self, | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/exporter.py", line 206, in def_func | |
| return_py_value = f(*argument_py_tree) | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/program_api.py", line 452, in invoke_with_self | |
| return func_def.callable(self, *args, **kwargs) | |
| File "/Users/y/w/iree-ios/iree-jax/models/gpt2/export.py", line 68, in encode | |
| kv, x = self._encode(self._params, prompt, t) | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/tracing.py", line 53, in __call__ | |
| return current_ir_trace().handle_call(self, args, kwargs) | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/tracing.py", line 113, in handle_call | |
| return target.resolve_call(self, *args, **kwargs) | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/builtins.py", line 60, in resolve_call | |
| lowered = self.jit_f.lower(*abstract_args) | |
| File "/Users/y/miniforge3/envs/transformers/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 163, in reraise_with_filtered_traceback | |
| return fun(*args, **kwargs) | |
| File "/Users/y/miniforge3/envs/transformers/lib/python3.10/site-packages/jax/_src/pjit.py", line 376, in lower | |
| donate_argnums) = infer_params_fn(*args, **kwargs) | |
| File "/Users/y/miniforge3/envs/transformers/lib/python3.10/site-packages/jax/_src/api.py", line 443, in infer_params | |
| return pjit.common_infer_params(pjit_info_args, *args, **kwargs) | |
| File "/Users/y/miniforge3/envs/transformers/lib/python3.10/site-packages/jax/_src/pjit.py", line 483, in common_infer_params | |
| in_shardings = out_shardings = _create_sharding_with_device_backend( | |
| File "/Users/y/miniforge3/envs/transformers/lib/python3.10/site-packages/jax/_src/pjit.py", line 781, in _create_sharding_with_device_backend | |
| xb.get_backend(backend).get_default_device_assignment(1)[0]) | |
| File "/Users/y/miniforge3/envs/transformers/lib/python3.10/site-packages/jax/_src/lib/xla_bridge.py", line 493, in get_backend | |
| return _get_backend_uncached(platform) | |
| File "/Users/y/miniforge3/envs/transformers/lib/python3.10/site-packages/jax/_src/lib/xla_bridge.py", line 485, in _get_backend_uncached | |
| raise RuntimeError(f"Unknown backend {platform}") | |
| jax._src.traceback_util.UnfilteredStackTrace: RuntimeError: Unknown backend iree | |
| The stack trace below excludes JAX-internal frames. | |
| The preceding is the original exception that occurred, unmodified. | |
| -------------------- | |
| The above exception was the direct cause of the following exception: | |
| Traceback (most recent call last): | |
| File "/Users/y/w/iree-ios/iree-jax/models/gpt2/export.py", line 110, in <module> | |
| absl.app.run(main) | |
| File "/Users/y/miniforge3/envs/transformers/lib/python3.10/site-packages/absl/app.py", line 308, in run | |
| _run_main(main, args) | |
| File "/Users/y/miniforge3/envs/transformers/lib/python3.10/site-packages/absl/app.py", line 254, in _run_main | |
| sys.exit(main(argv)) | |
| File "/Users/y/w/iree-ios/iree-jax/models/gpt2/export.py", line 101, in main | |
| f.write(str(Program.get_mlir_module(module))) | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/program_api.py", line 377, in get_mlir_module | |
| info = Program.get_info(Program._get_instance(m)) | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/program_api.py", line 372, in _get_instance | |
| m = m() | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/program_api.py", line 459, in __new__ | |
| export_function() | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/program_api.py", line 454, in export_function | |
| info.export_module.def_func(invoke_with_self, | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/exporter.py", line 206, in def_func | |
| return_py_value = f(*argument_py_tree) | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/program_api.py", line 452, in invoke_with_self | |
| return func_def.callable(self, *args, **kwargs) | |
| File "/Users/y/w/iree-ios/iree-jax/models/gpt2/export.py", line 68, in encode | |
| kv, x = self._encode(self._params, prompt, t) | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/tracing.py", line 53, in __call__ | |
| return current_ir_trace().handle_call(self, args, kwargs) | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/tracing.py", line 113, in handle_call | |
| return target.resolve_call(self, *args, **kwargs) | |
| File "/Users/y/w/iree-ios/iree-jax/iree/jax/builtins.py", line 60, in resolve_call | |
| lowered = self.jit_f.lower(*abstract_args) | |
| RuntimeError: Unknown backend iree |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment