Skip to content

Commit 6123fdd

Browse files
committed
docs: resolve leftover lora_llama3_demo.ipynb comments and relative code links in lora_model_bringup.md
1 parent c5bc2e4 commit 6123fdd

2 files changed

Lines changed: 9 additions & 15 deletions

File tree

docs/guides/lora_model_bringup.md

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,12 +30,12 @@ To enable LoRA support for a new model, follow these two simple steps:
3030

3131
The target model architecture must already be implemented and supported as a base model in MaxText.
3232

33-
- The JAX/NNX model definition should be located under `src/maxtext/models/` (e.g., \[gemma3.py\](../../src/maxtext/models/gemma3.py)).
33+
- The JAX/NNX model definition should be located under `src/maxtext/models/` (e.g., [gemma3.py](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/gemma3.py)).
3434
- The model configurations must be registered and runnable for baseline pre-training or full fine-tuning.
3535

3636
### Step 1.2: Define Trainable LoRA Target Modules
3737

38-
Add a recommended target pattern for your model architecture prefix in \[src/maxtext/configs/post_train/lora_module_path.yml\](../../src/maxtext/configs/post_train/lora_module_path.yml):
38+
Add a recommended target pattern for your model architecture prefix in [src/maxtext/configs/post_train/lora_module_path.yml](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/configs/post_train/lora_module_path.yml):
3939

4040
```yaml
4141
your_model_prefix: "decoder/layers/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))"
@@ -69,7 +69,7 @@ If you want to perform decoding or run high-performance serving on your adapted
6969
To add weight mapping for vLLM decode:
7070

7171
1. **Create a Weight Mapping Config**:
72-
Create a new file in \[src/maxtext/integration/tunix/weight_mapping/\](../../src/maxtext/integration/tunix/weight_mapping/) (e.g., `your_model.py`) defining a mapping dataclass. You can refer to \[gemma3.py\](../../src/maxtext/integration/tunix/weight_mapping/gemma3.py) or \[llama3.py\](../../src/maxtext/integration/tunix/weight_mapping/llama3.py) as templates.
72+
Create a new file in [src/maxtext/integration/tunix/weight_mapping/](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/integration/tunix/weight_mapping/) (e.g., `your_model.py`) defining a mapping dataclass. You can refer to [llama3.py](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/integration/tunix/weight_mapping/llama3.py) as a template.
7373

7474
Your class should specify:
7575

@@ -78,7 +78,7 @@ To add weight mapping for vLLM decode:
7878
- `lora_to_hf_mappings()`: Custom mapping for LoRA weights if they require different handling.
7979

8080
2. **Register the Mapping**:
81-
Register your new class in \[src/maxtext/integration/tunix/weight_mapping/__init__.py\](../../src/maxtext/integration/tunix/weight_mapping/__init__.py) inside the `StandaloneVllmWeightMapping` class:
81+
Register your new class in [src/maxtext/integration/tunix/weight_mapping/__init__.py](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/integration/tunix/weight_mapping/__init__.py) inside the `StandaloneVllmWeightMapping` class:
8282

8383
```python
8484
# Inside StandaloneVllmWeightMapping

src/maxtext/examples/lora_llama3_demo.ipynb

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,10 @@
7979
" # Install uv, a fast Python package installer\n",
8080
" !pip install uv\n",
8181
" \n",
82+
" # Set the torch backend to CPU for uv\n",
83+
" import os\n",
84+
" os.environ[\"UV_TORCH_BACKEND\"] = \"cpu\"\n",
85+
" \n",
8286
" # Install MaxText and post-training dependencies\n",
8387
" !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
8488
" !install_tpu_post_train_extra_deps"
@@ -215,16 +219,6 @@
215219
"outputs": [],
216220
"source": [
217221
"if not epath.Path(MODEL_CHECKPOINT_PATH).exists():\n",
218-
" # Install torch for the conversion script\n",
219-
" print(\"Installing torch...\")\n",
220-
" subprocess.run(\n",
221-
" [\n",
222-
" sys.executable, \"-m\", \"pip\", \"install\",\n",
223-
" \"torch\", \"--index-url\", \"https://download.pytorch.org/whl/cpu\"\n",
224-
" ],\n",
225-
" check=True\n",
226-
" )\n",
227-
"\n",
228222
" print(\"Converting checkpoint from HuggingFace...\")\n",
229223
" env = os.environ.copy()\n",
230224
" env[\"JAX_PLATFORMS\"] = \"cpu\"\n",
@@ -607,4 +601,4 @@
607601
},
608602
"nbformat": 4,
609603
"nbformat_minor": 5
610-
}
604+
}

0 commit comments

Comments
 (0)