Skip to content

Add sentence configure memory allocator in docs #168

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wesselb merged 1 commit into from
Nov 18, 2025
Merged

Add sentence configure memory allocator in docs #168

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions docs/finetuning.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ model = AuroraPretrained()
model.load_checkpoint()
```

## Basic Fine-Tuning Environment
## Fine-Tuning Environment

We provide a very basic Docker image and fine-tuning loop to get you started.
This Docker image is built from a NVIDIA PyTorch base image,
Expand All @@ -30,10 +30,13 @@ docker run --rm -it -v .:/app/aurora \
Then, within the image, execute

```bash
python finetuning/finetune.py
PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync \
python finetuning/finetune.py
```

to run the sample fine-tuning loop.
`PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync` enables CUDA's built-in
asynchronous memory allocator, which is recommended for Aurora.
This loop should run on an A100 with 80 GB of memory.
If you need to reduce memory usage, you could try the following:
(a) split the model and optimiser parameters across multiple GPUs with
Expand Down
Loading