berkus
/
ANE
mirror of https://github.com/maderix/ANE.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add --data path support for static training pipelines

This commit is contained in:
nabbilkhan 2026-03-03 19:19:49 +00:00
parent 443194bca4
commit c04168ee17
3 changed files with 60 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
training/README.md
View File

@ -83,15 +83,17 @@ Downloads pretokenized TinyStories (Llama 2 BPE, 32K vocab) from HuggingFace. Pr
### 2. Build & Train
```bash
# Static baseline (classifier + softmax on CPU)
make train_large
./train_large stories110M.bin 256 100 1e-4
./train_large --model stories110M.bin --steps 100 --lr 1e-4
# PR#19: ANE-offloaded classifier + softmax + rmsnorm_bwd
make train_large_ane
./train_large_ane stories110M.bin 256 100 1e-4
./train_large_ane --no-ane-extras --steps 100 # disable ANE extras
# Static baseline (classifier + softmax on CPU)
make train_large
./train_large stories110M.bin 256 100 1e-4
./train_large --model stories110M.bin --steps 100 --lr 1e-4
./train_large --data ./tinystories_data00.bin --steps 100 --lr 1e-4
# PR#19: ANE-offloaded classifier + softmax + rmsnorm_bwd
make train_large_ane
./train_large_ane stories110M.bin 256 100 1e-4
./train_large_ane --no-ane-extras --steps 100 # disable ANE extras
./train_large_ane --data ./tinystories_data00.bin --steps 100 --lr 1e-4
# Dynamic pipeline (no recompilation)
cd training_dynamic && make train
@ -100,13 +102,14 @@ cd training_dynamic && make train
./train --steps 200 --lr 1e-4 # custom steps/lr
```
**CLI flags (all pipelines):**
**CLI flags (`train_large` / `train_large_ane`):**
- `--steps N` (default 10000)
- `--lr F` (default 3e-4)
- `--model PATH` — pretrained weights file
- `--ckpt PATH` — checkpoint file (preserved across exec() restarts)
- `--resume` — resume from checkpoint
- `--no-ane-extras` — (train_large_ane only) disable ANE classifier/softmax/rmsnorm_bwd
- `--lr F` (default 3e-4)
- `--model PATH` — pretrained weights file
- `--data PATH` — tokenized TinyStories `.bin` file (default: `tinystories_data00.bin`)
- `--ckpt PATH` — checkpoint file (preserved across exec() restarts)
- `--resume` — resume from checkpoint
- `--no-ane-extras` — (train_large_ane only) disable ANE classifier/softmax/rmsnorm_bwd
### 3. Monitor with Dashboard

54
training/train_large.m
View File

@ -5,9 +5,9 @@
#include "stories_mil.h"
#include "stories_cpu_ops.h"
#define CKPT_PATH_DEFAULT "ane_stories110M_ckpt.bin"
#define MODEL_PATH_DEFAULT "stories110M.bin"
#define DATA_PATH "tinystories_data00.bin"
#define CKPT_PATH_DEFAULT "ane_stories110M_ckpt.bin"
#define MODEL_PATH_DEFAULT "stories110M.bin"
#define DATA_PATH_DEFAULT "tinystories_data00.bin"
// ===== Weight loading from llama2.c format =====
static bool load_pretrained(LayerWeights *lw, float *rms_final, float *embed, const char *path) {
@ -192,22 +192,24 @@ int main(int argc, char *argv[]) {
float adam_b1=0.9f, adam_b2=0.999f, adam_eps=1e-8f;
int adam_t = 0, start_step = 0;
// Parse args
const char *ckpt_path = CKPT_PATH_DEFAULT;
const char *model_path = MODEL_PATH_DEFAULT;
bool do_resume = false;
int pos = 0;
for (int i=1; i<argc; i++) {
if (strcmp(argv[i], "--resume") == 0) do_resume = true;
else if (strcmp(argv[i], "--steps") == 0 && i+1<argc) total_steps = atoi(argv[++i]);
else if (strcmp(argv[i], "--lr") == 0 && i+1<argc) lr = atof(argv[++i]);
else if (strcmp(argv[i], "--ckpt") == 0 && i+1<argc) ckpt_path = argv[++i];
else if (strcmp(argv[i], "--model") == 0 && i+1<argc) model_path = argv[++i];
else if (argv[i][0] != '-') {
if (pos == 0) model_path = argv[i];
else if (pos == 1) { /* seq - compile-time constant */ }
else if (pos == 2) total_steps = atoi(argv[i]);
else if (pos == 3) lr = atof(argv[i]);
// Parse args
const char *ckpt_path = CKPT_PATH_DEFAULT;
const char *model_path = MODEL_PATH_DEFAULT;
const char *data_path = DATA_PATH_DEFAULT;
bool do_resume = false;
int pos = 0;
for (int i=1; i<argc; i++) {
if (strcmp(argv[i], "--resume") == 0) do_resume = true;
else if (strcmp(argv[i], "--steps") == 0 && i+1<argc) total_steps = atoi(argv[++i]);
else if (strcmp(argv[i], "--lr") == 0 && i+1<argc) lr = atof(argv[++i]);
else if (strcmp(argv[i], "--ckpt") == 0 && i+1<argc) ckpt_path = argv[++i];
else if (strcmp(argv[i], "--model") == 0 && i+1<argc) model_path = argv[++i];
else if (strcmp(argv[i], "--data") == 0 && i+1<argc) data_path = argv[++i];
else if (argv[i][0] != '-') {
if (pos == 0) model_path = argv[i];
else if (pos == 1) { /* seq - compile-time constant */ }
else if (pos == 2) total_steps = atoi(argv[i]);
else if (pos == 3) lr = atof(argv[i]);
pos++;
}
}
@ -283,8 +285,12 @@ int main(int argc, char *argv[]) {
}
// mmap token data
int data_fd = open(DATA_PATH, O_RDONLY);
if (data_fd < 0) { printf("Cannot open %s\n", DATA_PATH); return 1; }
int data_fd = open(data_path, O_RDONLY);
if (data_fd < 0) {
printf("Cannot open token data: %s\n", data_path);
printf("Hint: run `bash download_data.sh` in training/ or pass --data /path/to/tinystories_data00.bin\n");
return 1;
}
struct stat st; fstat(data_fd, &st);
size_t data_len = st.st_size;
uint16_t *token_data = (uint16_t*)mmap(NULL, data_len, PROT_READ, MAP_PRIVATE, data_fd, 0);
@ -340,9 +346,9 @@ int main(int argc, char *argv[]) {
lw, la, rms_final, &arms_final, embed, &aembed);
printf("[exec() restart step %d, %d compiles, loss=%.4f]\n", step, g_compile_count, last_loss);
fflush(stdout);
execl(argv[0], argv[0], "--resume", "--ckpt", ckpt_path, NULL);
perror("execl"); return 1;
}
execl(argv[0], argv[0], "--resume", "--ckpt", ckpt_path, "--data", data_path, NULL);
perror("execl"); return 1;
}
// Compile all layers' weight-bearing kernels
uint64_t tc = mach_absolute_time();

18
training/train_large_ane.m
View File

@ -9,7 +9,7 @@
// NLL loss + gradient (needs target indexing)
//
// Build: make train_large_ane
// Run: ./train_large_ane [--resume] [--steps N] [--lr F]
// Run: ./train_large_ane [--resume] [--steps N] [--lr F] [--data PATH]
#include "stories_io.h"
#include "stories_mil.h"
#include "stories_cpu_ops.h"
@ -18,7 +18,7 @@
#define CKPT_PATH_DEFAULT "ane_stories110M_ckpt.bin"
#define MODEL_PATH_DEFAULT "stories110M.bin"
#define DATA_PATH "tinystories_data00.bin"
#define DATA_PATH_DEFAULT "tinystories_data00.bin"
// ===== Weight loading from llama2.c format =====
static bool load_pretrained(LayerWeights *lw, float *rms_final, float *embed, const char *path) {
@ -204,6 +204,7 @@ int main(int argc, char *argv[]) {
int adam_t = 0, start_step = 0;
const char *ckpt_path = CKPT_PATH_DEFAULT;
const char *model_path = MODEL_PATH_DEFAULT;
const char *data_path = DATA_PATH_DEFAULT;
bool do_resume = false;
bool ane_extras = true; // classifier, softmax, rmsnorm_bwd on ANE
int pos = 0;
@ -214,6 +215,7 @@ int main(int argc, char *argv[]) {
else if (strcmp(argv[i], "--lr") == 0 && i+1<argc) lr = atof(argv[++i]);
else if (strcmp(argv[i], "--ckpt") == 0 && i+1<argc) ckpt_path = argv[++i];
else if (strcmp(argv[i], "--model") == 0 && i+1<argc) model_path = argv[++i];
else if (strcmp(argv[i], "--data") == 0 && i+1<argc) data_path = argv[++i];
else if (argv[i][0] != '-') {
if (pos == 0) model_path = argv[i];
else if (pos == 1) { /* seq - compile-time constant */ }
@ -271,8 +273,12 @@ int main(int argc, char *argv[]) {
}
// mmap token data
int data_fd = open(DATA_PATH, O_RDONLY);
if (data_fd < 0) { printf("Cannot open %s\n", DATA_PATH); return 1; }
int data_fd = open(data_path, O_RDONLY);
if (data_fd < 0) {
printf("Cannot open token data: %s\n", data_path);
printf("Hint: run `bash download_data.sh` in training/ or pass --data /path/to/tinystories_data00.bin\n");
return 1;
}
struct stat st; fstat(data_fd, &st);
size_t data_len = st.st_size;
uint16_t *token_data = (uint16_t*)mmap(NULL, data_len, PROT_READ, MAP_PRIVATE, data_fd, 0);
@ -354,9 +360,9 @@ int main(int argc, char *argv[]) {
printf("[exec() restart step %d, %d compiles, loss=%.4f]\n", step, g_compile_count, last_loss);
fflush(stdout);
if (ane_extras)
execl(argv[0], argv[0], "--resume", "--ckpt", ckpt_path, NULL);
execl(argv[0], argv[0], "--resume", "--ckpt", ckpt_path, "--data", data_path, NULL);
else
execl(argv[0], argv[0], "--resume", "--ckpt", ckpt_path, "--no-ane-extras", NULL);
execl(argv[0], argv[0], "--resume", "--ckpt", ckpt_path, "--data", data_path, "--no-ane-extras", NULL);
perror("execl"); return 1;
}