google / google/tunix

[Bug] Fix run_prod for scripts/grpo_demo_llama3_qwen2.py with sglang_jax with LoRA

Open
#995 1 comment 0 reactions 1 assignee Claimed by @aolemila View on GitHub
type:bug
Dominant language
Python
Stars
2.5k
Forks
345
Avg merge
1d 7h
Merged PRs (30d)
240

Description

**Expected Behavior**

Pass.

**Actual Behavior**

Failed with the following error:

```bash
File "/usr/local/lib/python3.12/site-packages/jax/_src/tree_util.py", line 361, in tree_map
return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/jax/_src/tree_util.py", line 361, in
return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
^^^^^^
File "/usr/local/lib/python3.12/site-packages/qwix/_src/flax_util.py", line 293, in fn
return unbox(x.get_raw_value())
^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/flax/nnx/variablelib.py", line 299, in __getattr__
return getattr(self.raw_value, name)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'jaxlib._jax.ArrayImpl' object has no attribute 'get_raw_value'
One or more scripts failed!
```
See more [here](https://github.com/google/tunix/actions/runs/21205448021/job/61000751630#step:7:11846).

**Steps to Reproduce the Problem**

```bash
unset JAX_PLATFORMS
pip list | egrep 'jax|flax|libtpu'
cd ..
git clone https://github.com/sgl-project/sglang-jax.git && cd sglang-jax/python && pip install -e . && cd ../..
pip install jax==0.8.1 flax==0.12.0 libtpu==0.0.24
pip list | egrep 'jax|flax|libtpu'
cd tunix
```

Run `python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=20 --rollout-engine sglang_jax --enable-lora --lora-target-modules all`.

**Environment**

- **OS:** [e.g., Ubuntu, etc.]
- **Project Version:** [e.g., 0.0.1]

**Checklist**

- [ ] I have searched the existing issues for a similar bug report.
- [ ] I have provided all the required information in the "Environment" section.
- [ ] I have provided a minimal, reproducible example.

**Would you like to help us fix it?**

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.