Skip to content
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.

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

Add training callback to send predictions to WandB table #521

Merged
merged 16 commits into from
Sep 13, 2023

Conversation

Glavin001
Copy link
Contributor

@Glavin001 Glavin001 commented Sep 3, 2023

How to use: #521 (comment)


Closes #490

What's New?

See examples:
https://www.loom.com/share/acaa23516b524aa29328b87b90f82599?sid=e2796761-d398-46a8-96a9-7f965b77437c

accelerate launch scripts/finetune.py examples/llama-2/tiny-random.yml

How to configure

1️⃣ Enable WandB
2️⃣ Every Eval will push updates to WandB

  • Change eval_steps in config as desired

Tasks

  • Table should include:
    • Prompt
    • [WIP] Predicted completion
    • Expected completion
    • Score / evaluation
  • Record for each evaluation
    • Configurable?
  • Ability to compare between 2 steps in WandB
  • Clean up old code, remove debugging code, etc

@Glavin001
Copy link
Contributor Author

Making progress, very slowly 😆 Learning a lot more about Python, Pytorch, HF Transformers though, which is good!

image

src/axolotl/utils/callbacks.py Outdated Show resolved Hide resolved
src/axolotl/utils/callbacks.py Outdated Show resolved Hide resolved
src/axolotl/utils/callbacks.py Outdated Show resolved Hide resolved
src/axolotl/utils/callbacks.py Outdated Show resolved Hide resolved
src/axolotl/utils/callbacks.py Outdated Show resolved Hide resolved
src/axolotl/utils/callbacks.py Outdated Show resolved Hide resolved
src/axolotl/utils/trainer.py Outdated Show resolved Hide resolved
src/axolotl/utils/callbacks.py Outdated Show resolved Hide resolved
@Glavin001
Copy link
Contributor Author

Today's progress update: https://www.loom.com/share/1d1cb34d846440e2a5258fd3593a9e80

@Glavin001
Copy link
Contributor Author

docker-compose.yaml Outdated Show resolved Hide resolved
scripts/finetune.py Outdated Show resolved Hide resolved
@Glavin001 Glavin001 changed the title [WIP] Add training callback to send predictions to WandB table Add training callback to send predictions to WandB table Sep 9, 2023
@Glavin001 Glavin001 marked this pull request as ready for review September 9, 2023 07:53
src/axolotl/utils/trainer.py Outdated Show resolved Hide resolved
src/axolotl/utils/callbacks.py Outdated Show resolved Hide resolved
.vscode/launch.json Outdated Show resolved Hide resolved
prompt_encoding = tokenizer(
prompt_texts, padding=True, return_tensors="pt"
).to(self.cfg.device)
predictions = trainer.model.generate(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

trainer.prediction_step(...) might be easier to use

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It was. However, I had that previously and it appeared to output strange predictions. At one point I actually had both to compare and only model.generate was useful.

I'm struggling to find the WandB report. I'll try again and see how it goes.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Top is trainer.prediction_step and bottom is trainer.model.generate:

image

Not sure if this is correct though? I originally showed both, I could add it too?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added both now

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

image

@Glavin001 Glavin001 mentioned this pull request Sep 12, 2023
4 tasks
src/axolotl/utils/trainer.py Show resolved Hide resolved
src/axolotl/utils/config.py Show resolved Hide resolved
@winglian
Copy link
Collaborator

also, the pre-commit checks are failing, just run pre-commit run --all-files

Copy link
Collaborator

@winglian winglian left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks!

@Glavin001
Copy link
Contributor Author

PS. Whenever you merge please squash the commits. Normally I’d have nice atomic/semantically sensible commits, although this PR had a ton of quick trial and error and saving WIP on remote GPU servers smile😁

@winglian winglian merged commit 5b67ea9 into axolotl-ai-cloud:main Sep 13, 2023
3 checks passed
@teknium1
Copy link
Contributor

PS. Whenever you merge please squash the commits. Normally I’d have nice atomic/semantically sensible commits, although this PR had a ton of quick trial and error and saving WIP on remote GPU servers smile😁

So how do I enable this and where do I set the prompts I want it to inference

@Glavin001
Copy link
Contributor Author

The new options (copied from the README, section Config > All yaml options, docs could be improved) are:

eval_table_size: # Approximate number of predictions sent to wandb depending on batch size. Enabled above 0. Default is 0
eval_table_max_new_tokens: # Total number of tokens generated for predictions sent to wandb. Default is 128

For example,

eval_table_size: 5
eval_table_max_new_tokens: 64

will:

  • Enable the predictions table being sent to WandB.
  • Have a table with at least 5 examples (depends on batch size, if there are 4 within a batch you may get 8 instead, Math.ceil(eval_table_size / batch_size) * batch_size).
  • Each prediction in table will generate at most 64 new tokens in length

The examples are currently extracted from the eval dataset, which is automatically set aside via the val_set_size option:

# How much of the dataset to set aside as evaluation. 1 = 100%, 0.50 = 50%, etc. 0 for no eval.
val_set_size: 0.04

In the future, a separate dataset entirely could be used, providing more control.
Still more work to be done on this feature.

Hope this helps! Let us know how it goes testing it, thanks!

@teknium1
Copy link
Contributor

teknium1 commented Oct 22, 2023

The new options (copied from the README, section Config > All yaml options, docs could be improved) are:

eval_table_size: # Approximate number of predictions sent to wandb depending on batch size. Enabled above 0. Default is 0
eval_table_max_new_tokens: # Total number of tokens generated for predictions sent to wandb. Default is 128

For example,

eval_table_size: 5
eval_table_max_new_tokens: 64

will:

  • Enable the predictions table being sent to WandB.
  • Have a table with at least 5 examples (depends on batch size, if there are 4 within a batch you may get 8 instead, Math.ceil(eval_table_size / batch_size) * batch_size).
  • Each prediction in table will generate at most 64 new tokens in length

The examples are currently extracted from the eval dataset, which is automatically set aside via the val_set_size option:

# How much of the dataset to set aside as evaluation. 1 = 100%, 0.50 = 50%, etc. 0 for no eval.
val_set_size: 0.04

In the future, a separate dataset entirely could be used, providing more control. Still more work to be done on this feature.

Hope this helps! Let us know how it goes testing it, thanks!

I see, ok. I am hoping to have a seperate dataset of prompts mainly because afaik axolotl does not allow setting a seperate eval dataset, it just comes randomly from % of the dataset to train on, and I have a very specific set of prompts not in the training data I need to eval with, I also dont use normal evals at all, because I have a very specific dataset that needs every entry trained on, and I dont trust eval loss anyways, so I'd rather not section some % of my dataset for evals

mkeoliya pushed a commit to mkeoliya/axolotl that referenced this pull request Dec 15, 2023
…cloud#521)

* WIP Add training callback to send predictions to WandB table

* WIP improve wandb table reporting callback

* WIP improve wandb table reporting callback (cont)

* Add VSCode launching for debugging

* Add tiny llama example

* WIP attempt to improve post-eval prediction generation for table

* WIP attempt to improve post-eval prediction generation for table - part 2

* WIP batch generation

* WIP attempt to handle sample_packing using position_ids for wandb prediction table

* WIP add code for debugging

* Fix sample_packing support for wandb prediction table

* Clean up code for PR review

* Add eval_table_size, eval_table_max_new_tokens configs & clean up code

* Clean up PR, delete VSCode config, add tiny-llama example

* Add eval_table_size, eval_table_max_new_tokens documentation. Fix linting/formatting
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Add prediction (table) artifacts to Weights & Biases logger
3 participants