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.

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

Sync internal repo to external Dec 12 2024 #103

Open
wants to merge 139 commits into
base: main
Choose a base branch
from
Open

Sync internal repo to external Dec 12 2024 #103

wants to merge 139 commits into from

Conversation

hannanjgaws
Copy link
Contributor

Issue #, if available:

Description of changes:

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

aws-patlange and others added 30 commits December 12, 2024 21:59
@aws-patlange @hannanjgaws
[generation] Fix on-device generation without filtering
2e807a6 
GitOrigin-RevId: a90252cdbff3cbf2fd53211b5df54a93499e7ef3
@hannanjgaws
[sos] mulit-bucket and bsh attn layout support
3d965f6 
GitOrigin-RevId: 1cae3bb498801f764e33126cba9ace01904ddbfc
@hannanjgaws
[generation_demo] fixes, add flags
1e0bbe1 
GitOrigin-RevId: f003d9b12bbb2d3dad44e985c575f184fecbb392
@hannanjgaws
[sos] seq mismatch fix for testing
d7725a7 
 - pad n_positions for flash decoding

GitOrigin-RevId: a18dbafde388c95b2584d13324bdaa6a865355dd
@aws-amerrez @hannanjgaws
sequence parallel with fully unrolled and on_device embedding
18f5eab 
GitOrigin-RevId: 7742ae57c97e499bf3d3d76163e8b1d24b952f4f
@jluntamazon @hannanjgaws
[compiler] Enabled writing HLO file metadata when dumping is configured
f080045 
GitOrigin-RevId: fe2b7dfb9f19de7af43822606b1600f67d8adb70
@hannanjgaws
[hlo] switch to the fast past cumsum NKI kernel for floats
bda234f 
GitOrigin-RevId: 7e3a08ef18016931b541054428bc5209aae78a71
@hannanjgaws
[speculation] Fix 2D cache_id pad for speculative_forward
4365742 
GitOrigin-RevId: cf907eeaf917184f1e96ffe9e62b6d9b2ff93ab5
@hannanjgaws
Enabling output bucketing for continuous batching with SBH cache layout.
d578c19 
GitOrigin-RevId: 7d735907aad75c99c43c65f59130da51036d3158
@hannanjgaws
[config/attention/llama] Transposed attention output projection weight
6b7b980 
GitOrigin-RevId: 133a97962ed104ae092fba6e43413f99de01e7b4
@aws-bowencc @hannanjgaws
make weightiling permute order configurable
0fa7550 
GitOrigin-RevId: eb9bfaa7eb27023355b2312559aa4f43a82e55e4
@aws-bowencc @hannanjgaws
make weightiling permute order configurable
f9bc7c7 
GitOrigin-RevId: 4872e762d44756433ba3da92ef989712aa21619c
@hannanjgaws
[attention] Add support for masking exp values in context
6e2b374 
GitOrigin-RevId: 2a353ea9548af30f102b75e71b673e699bccdd26
@aws-patlange @hannanjgaws
[generation_demo] Add naive continuous batching
c038af9 
GitOrigin-RevId: 3975a0b10c329fbdcaf37d46e4d940704d3ac1bc
@gevadyla @hannanjgaws
switch back to gelu_new_legacy
008ed9a 
GitOrigin-RevId: 86c4a6eb93c287141fcc31eeb65c804e53e72e8e
@aws-amulyaab @hannanjgaws
Enable continuous batching for multi-layer models
f7316ff 
GitOrigin-RevId: 883608c186597ba01eeae16b8251c29788f4e76a
@aws-patlange @hannanjgaws
Fix on-device generation during continuous batching
78e39e4 
GitOrigin-RevId: 122b82337d25abbd305eac0c6e4375af25992ca0
@aws-patlange @hannanjgaws
Fix on-device embedding and continuous batching
695b5d6 
GitOrigin-RevId: bdb3278f994a3b5a53d231d8150d9ae3962ed298
@aws-patlange @hannanjgaws
Use collectives_layout=BSH as default and add deprecation warning
3913740 
GitOrigin-RevId: 33b62dc57a682ed64fa58ab7df9d9b90fef6d163
@akhil-aws @hannanjgaws
Remove internal code name
478714e 
GitOrigin-RevId: f9787c6cd165eece8793f8c3032a3f7d754c8487
@aws-amerrez @hannanjgaws
stop criteria fix to cover generated tokens
4a0f913 
GitOrigin-RevId: 21956bf24d98634c339b3c4b5af370f90b98b290
@hannanjgaws
[attn] Fix NKI kernel import in attention utils.
cca53a1 
GitOrigin-RevId: c2341eda41d0e3d21e143df586de8ac196ad0312
@aws-patlange @hannanjgaws
Revert "stop criteria fix to cover generated tokens"
8ad9cc4 
This reverts commit 21956bf24d98634c339b3c4b5af370f90b98b290.

GitOrigin-RevId: a3694e9f2d0313ea4baaaac7c59494d4a4576a22
@aws-patlange @hannanjgaws
Revert "[generation_demo] Add naive continuous batching"
c0c54cc 
This reverts commit 3975a0b10c329fbdcaf37d46e4d940704d3ac1bc.

GitOrigin-RevId: e24f95c9454f83129ae6ebcbb71978113c84f4fc
@hannanjgaws
[sos] node interleave qkv support
1ba1ba7 
GitOrigin-RevId: 77686e5f269124b82348d1bda5df64e8a16fc6eb
@aws-patlange @hannanjgaws
Revert "Use collectives_layout=BSH as default and add deprecation war…
86a6dd9 
…ning"

This reverts commit 33b62dc57a682ed64fa58ab7df9d9b90fef6d163.

GitOrigin-RevId: 53bb56e4f9d36e0aaeb876f9d916954d9aa92bb2
@aws-yishanm @hannanjgaws
Allow saving and loading of preshareded weights
e9a0159 
GitOrigin-RevId: 8300b76c09c017c9034c4dbd8685a02fa5c52538
@hannanjgaws
Extend sequence parallel normalization support in the MLP and LM Head…
7fb318c 
… functions

GitOrigin-RevId: fb5f0215163939ca54bf3390efed127b4fd29c74
@hannanjgaws
[sos] unit_2 bug fix
ef6b434 
GitOrigin-RevId: 78f4fa47ee75e45bfc15c40c7eb3bf944449438c
@hannanjgaws
[decoder] Fix string comparison for 'None' params
4f829a0 
GitOrigin-RevId: 201bd7397e146fa163c768c98ed850245ad02687
Yi-Hsiang (Sean) Lai and others added 25 commits December 12, 2024 22:04
@hannanjgaws
fix greedy generation
e05fc47 
GitOrigin-RevId: a5b4be50f7f62254ce95b18afb347fb385a40bdb
@hannanjgaws
make sure we have the right masking first
9d27625 
GitOrigin-RevId: 667d84dc558133ebf1e3c601fbbd451306ee41ee
@hannanjgaws
Fix eagle accuracy issue
916e0cb 
GitOrigin-RevId: e3eac46a0d3c6cf55eb48d51ed4a809b047d3b68
@liangfu @hannanjgaws
Chunked Prefill
e03976c 
GitOrigin-RevId: 12736eb3efa948a5e4c9b39711d9b86ee9214806
@aws-bowencc @hannanjgaws
add option to skip allgather with duplicate q weights during sos
4248af8 
GitOrigin-RevId: a6f9bc02e7c5eb24875333b695edcc771ad05359
@hannanjgaws
Fix Eagle SD error due to additional inputs from chunked prefill
b3c588e 
GitOrigin-RevId: cfe4f155c95af5192aa1b289116ed48ae78dacd2
@hannanjgaws
Merge branch 'jf-sos-chunked-pa' into mainline
68007ff 
GitOrigin-RevId: 5b5442bc5045de324d305e9ab6cc07b6627367c3
@hannanjgaws
enable flashattention-style comm, duplicate query, topoaware-comm and
8b9112a 
in-memory-copy kv cache in shard over sequence for PagedAttention with
Chunked Prefill
Merge branch 'mainline' of trn2:git/KaenaTransformers into mainline

GitOrigin-RevId: 9b9cc41cb04f3e4fab2d5416d75fbd98b0b6f5ea
@hannanjgaws
Merge branch 'mainline' of ssh://git.amazon.com:2222/pkg/KaenaTransfo…
aad25fe 
…rmers into mainline

GitOrigin-RevId: d0e58424b388e14f836431e16bcf08cf2eb70170
@hannanjgaws
reorder argument input for shraded_kv_indexing
e861b55 
GitOrigin-RevId: 4457a4f55e760b7074a842a381d844c76ef73877
@hannanjgaws
fix softmax, topo-aware sharding for llama, optimize pre-layer mask g…
c4d0b2e 
…eneration

GitOrigin-RevId: 438c99a6e89299f043411115b292c12150678f9b
@hannanjgaws
remove topo-aware for shard-over-sequence
5aad935 
GitOrigin-RevId: 8bc51f70c4192779c90fb76efe164f2d1d6e085d
@hannanjgaws
Merge branch 'mainline' of ssh://git.amazon.com:2222/pkg/KaenaTransfo…
67d762b 
…rmers into mainline

GitOrigin-RevId: 9a10a3ea06add1f6f4469417831d7ef040e9ef38
@hannanjgaws
Merge branch 'mainline' of ssh://git.amazon.com:2222/pkg/KaenaTransfo…
4df0cb0 
…rmers into mainline

GitOrigin-RevId: ac9032ae7823fa1d8e6259e5740d4f173db5c256
@aws-bowencc @hannanjgaws
slice the kv cache correctly when using sos
dbeb8b8 
GitOrigin-RevId: 14367c262d1b0a62b9aeb719c88595037dd27433
@hannanjgaws
Fix SD token selection logic
2214b99 
GitOrigin-RevId: cc449cd726374524ef777e18b910a166cfc904c4
@aws-bowencc @hannanjgaws
fix eagle sd on corner cases when input length is 1 or 2
feef47c 
GitOrigin-RevId: eafbb83875892bf5d473c4950c3a0c75b21b714b
@hannanjgaws
Use f32 rng for SD token selection
5d0c91e 
GitOrigin-RevId: 210123e945a509d7087f478a322cbea7974b3542
@hannanjgaws
TP 128 draft fix
4a8eefa 
GitOrigin-RevId: b68155341b188362a4cb7137d2724fb461a9a50a
@aws-patlange @hannanjgaws
[generation] Fix per_batch_line sampling param lookup for CB
cedd4e7 
GitOrigin-RevId: 5039ec6b067d1431201563659bb05b9c90e79bb0
@aws-patlange @hannanjgaws
[generation] Enforce desired dtype for on_device_generation sampling …
76c569d 
…params

GitOrigin-RevId: 13cd995f8a2aff8a66a39cc7a828e0d05c67d39e
@aws-bowencc @hannanjgaws
fix duplicate q when q/tp!=1
38027d4 
GitOrigin-RevId: bc5163b8e710ac3076650a3e391d7cd4f6e4feca
@aws-patlange @hannanjgaws
[load_weights] Nullify materialized parameters
d2e795b 
GitOrigin-RevId: 67ad932a09d31b52c5b5f90a7ac29b1a6d5d3198
@hannanjgaws
Support bf16 rms norm and new neuron config bf16_rms_norm=True to ena…
a391dd9 
…ble it

cr https://code.amazon.com/reviews/CR-159646815

GitOrigin-RevId: 543b713424da4a14fb40f0387f89813867b322df
@aws-yishanm @hannanjgaws
enforce max context_length_estimate <= max n_positions
997bf9d 
GitOrigin-RevId: e34241e62bc8684b7a8014945c3ae3c60b4d5b74
gsnaws
gsnaws approved these changes Dec 13, 2024
View reviewed changes
Copy link

@gsnaws gsnaws left a comment

Choose a reason for hiding this comment

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

lgtm

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@gsnaws gsnaws gsnaws approved these changes

@aws-maens aws-maens Awaiting requested review from aws-maens aws-maens is a code owner

@musunita musunita Awaiting requested review from musunita musunita is a code owner

@jluntamazon jluntamazon Awaiting requested review from jluntamazon

At least 2 approving reviews are required to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.