Skip to content

(feat) add attention logits to model output, add attention soft_cap to vanilla attention; (fix) DP sharding of batch, update dtype of memory tracking interval #209

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

Jump to bottom
Open
dvruette wants to merge 88 commits into
base: main
Choose a base branch
from
Open

(feat) add attention logits to model output, add attention soft_cap to vanilla attention; (fix) DP sharding of batch, update dtype of memory tracking interval #209

dvruette wants to merge 88 commits into from

Conversation

@dvruette
Copy link
Contributor

@dvruette dvruette commented Aug 2, 2025

Collection of various improvements:

  • Attention logits in model output: useful for experiment tracking and debugging
  • Attention softcap: used by Gemma, now also supported by vanilla attention
  • DP sharding of batch: now correctly shards the batch along the DP axis. Depends on a breaking change in eformer ([fix] Apply with_sharding_constraint recursively to pytrees eformer#4).
  • Minor change to check_interval dtype for memory tracking.

dvruette added 30 commits July 31, 2025 17:38
@dvruette
implement attn softcap in vanilla attention
c00a57c
@dvruette
add attention logits to attention output
db3a35e
@dvruette
change check_interval type to int or float
40f9ec3
@dvruette
[fix] make sure sharding gets applied to every array in batch
f2cd8e5
@dvruette
prettier memory summary printing
f81e3e3
@dvruette
improve memory monitoring display
0ec9171
@dvruette
graceful handling of failing to save state
42c6764
@dvruette
Merge branch 'erfanzar:main' into main
a261472
@dvruette
keep track of dataloading time
f8bf6e4
@dvruette
implement metrics aggregation
c86da8e
@dvruette
make wandb tags configurable
3116af9
@dvruette
fix metrics aggregation
38d8bf7
@dvruette
annotate profiler steps
678a6fb
@dvruette
fix training step annotation
b8b5424
@dvruette
Merge branch 'main' of github.com:erfanzar/EasyDeL
8efe3e4
@dvruette
fix vanilla attention
6546900
@dvruette
add profiling based on env variable
17acf24
@dvruette
fix tpu_setup script for v6e
408e565
@dvruette
Merge branch 'erfanzar:main' into main
b4b1fbd
@dvruette
revert to_state function
bb13fa0
@dvruette
try fix for to_state function
b3663af
@dvruette
debugging to_state method
ecf61be
@dvruette
debugging to_state method
7655aa9
@dvruette
fix to_state method
a8dc4e1
@dvruette
add logs for starting/stopping profiler
390f020
@dvruette
disable profiling options
69403ea
@dvruette
minor logging change
ccb05fe
@dvruette
refactor wandb init to allow for arbitrary kwargs
ce11645
@dvruette
fix case where entity is provided as wandb kwarg
ddb9e7f
@dvruette
fix bug when resuming wandb run
1f13fd3
@dvruette
Merge branch 'main' of github.com:erfanzar/EasyDeL
a7ccc6e
@dvruette
add step skipping for resuming from checkpoint
916582e
@dvruette
install eformer from dvruette/eformer
4edfcfd
@dvruette
fix checkpointing
f187f79
@dvruette
add ray docker image to build
903b632
@dvruette
update ray image
fa10e6e
@dvruette
update ray
3c067f4
@dvruette
update lockfile
8ef48ce
@dvruette
Merge branch 'main' of github.com:erfanzar/EasyDeL into merge-upstream
827a355
@dvruette
remove old fix that is no longer necessary
d194cd8
@dvruette
update lockfile
d740552
@dvruette
add support for uploading profiler traces to gcs
9ccf74d
@dvruette
add orbax saving
c8ffb7a
@dvruette
fix orbax checkpointing
a0c04af
@dvruette
load from orbax checkpoint
b080e20
@dvruette
checkpointing fixes
acf93b2
@dvruette
more checkpointing safety
5708b7c
@dvruette
fix type hint
2bb8f83
@dvruette
various minor fixes and improvements to training loop/checkpointing
0904c88
@dvruette
remove barrier after checkpointing
b441979
@dvruette
update ray to latest version
85472a5
@dvruette
fix continuing from checkpoint, fix unnecessary comms
73eebea
@dvruette
minor changes
3840afc
@dvruette
fix easydel auto flags
38b2380
@dvruette
make training loop compatible with new data sampler
69f5617
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@dvruette