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 to head #771

Merged
merged 759 commits into from
Jul 15, 2024
Merged

sync to head #771

merged 759 commits into from
Jul 15, 2024
This pull request is big! We’re only showing the most recent 250 commits.

Commits on Apr 11, 2024

  1. Merge pull request #588 from google:vertex_flag 

    PiperOrigin-RevId: 623909103
    maxtext authors committed Apr 11, 2024
    Configuration menu
    Copy the full SHA
    993f931 View commit details
    Browse the repository at this point in the history

  2. Move tpu end-to-end test scripts to tpu folder

    @NinaCai
    NinaCai committed Apr 11, 2024
    Configuration menu
    Copy the full SHA
    d676993 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #575 from prrathi:main 

    PiperOrigin-RevId: 623946526
    maxtext authors committed Apr 11, 2024
    Configuration menu
    Copy the full SHA
    a80c435 View commit details
    Browse the repository at this point in the history

  4. Merge branch 'main' into nina/mv_tpu_end_to_end_scripts

    @NinaCai
    NinaCai committed Apr 11, 2024
    Configuration menu
    Copy the full SHA
    0cdcdd1 View commit details
    Browse the repository at this point in the history

  5. Merge pull request #589 from google:nina/mv_tpu_end_to_end_scripts 

    PiperOrigin-RevId: 623965280
    maxtext authors committed Apr 11, 2024
    Configuration menu
    Copy the full SHA
    28a3279 View commit details
    Browse the repository at this point in the history

Commits on Apr 12, 2024

  1. unify WORKDIR to /deps

    @michelle-yooh
    michelle-yooh committed Apr 12, 2024
    Configuration menu
    Copy the full SHA
    16a05c0 View commit details
    Browse the repository at this point in the history

  2. Merge branch 'main' into yiinho-prebuilt-te

    @chajath
    chajath committed Apr 12, 2024
    Configuration menu
    Copy the full SHA
    240e25d View commit details
    Browse the repository at this point in the history

  3. Share GCS path between Gemma-7b tests

    @A9isha
    A9isha committed Apr 12, 2024
    Configuration menu
    Copy the full SHA
    5b8a3c3 View commit details
    Browse the repository at this point in the history

  4. Merge pull request #583 from google:anisha-run-multicluster-test 

    PiperOrigin-RevId: 624288653
    maxtext authors committed Apr 12, 2024
    Configuration menu
    Copy the full SHA
    3bba7a6 View commit details
    Browse the repository at this point in the history

  5. Add README for llama2-7B

    @michelle-yooh
    michelle-yooh committed Apr 12, 2024
    Configuration menu
    Copy the full SHA
    efb6d9e View commit details
    Browse the repository at this point in the history

  6. Merge pull request #582 from google:yooh-llama2-7b-config-readme 

    PiperOrigin-RevId: 624312270
    maxtext authors committed Apr 12, 2024
    Configuration menu
    Copy the full SHA
    71b78a1 View commit details
    Browse the repository at this point in the history

  7. Merge pull request #587 from google:yooh/fix_wordir 

    PiperOrigin-RevId: 624312415
    maxtext authors committed Apr 12, 2024
    Configuration menu
    Copy the full SHA
    e947d62 View commit details
    Browse the repository at this point in the history

  8. Merge remote-tracking branch 'origin/main' into yiinho-prebuilt-te

    @chajath
    chajath committed Apr 12, 2024
    Configuration menu
    Copy the full SHA
    5d12cce View commit details
    Browse the repository at this point in the history

  9. Merge branch 'yiinho-prebuilt-te' of https://github.com/google/maxtext … 

    …into yiinho-prebuilt-te
    @chajath
    chajath committed Apr 12, 2024
    Configuration menu
    Copy the full SHA
    dd8fef5 View commit details
    Browse the repository at this point in the history

  10. Merge pull request #547 from google:yiinho-prebuilt-te 

    PiperOrigin-RevId: 624316618
    maxtext authors committed Apr 12, 2024
    Configuration menu
    Copy the full SHA
    ebd39aa View commit details
    Browse the repository at this point in the history

Commits on Apr 16, 2024

  1. adding script to fix the style and adding modified/fixed files with l… 

    …ine length 125
    @ssusie
    ssusie committed Apr 16, 2024
    Configuration menu
    Copy the full SHA
    9343eec View commit details
    Browse the repository at this point in the history

  2. Merge pull request #596 from google:ssusie-pylint-125 

    PiperOrigin-RevId: 625444223
    maxtext authors committed Apr 16, 2024
    Configuration menu
    Copy the full SHA
    3e316bd View commit details
    Browse the repository at this point in the history

Commits on Apr 17, 2024

  1. Move apt install from rto_setup.sh to setup.sh

    @tonyjohnchen
    tonyjohnchen committed Apr 17, 2024
    Configuration menu
    Copy the full SHA
    cee962b View commit details
    Browse the repository at this point in the history

  2. Merge pull request #598 from google:rto 

    PiperOrigin-RevId: 625724059
    maxtext authors committed Apr 17, 2024
    Configuration menu
    Copy the full SHA
    de9715d View commit details
    Browse the repository at this point in the history

Commits on Apr 18, 2024

  1. Update instructions for installing snap.

    @RoshaniN
    RoshaniN committed Apr 18, 2024
    Configuration menu
    Copy the full SHA
    57c9905 View commit details
    Browse the repository at this point in the history

  2. Merge pull request #601 from google:RoshaniN-patch-1 

    PiperOrigin-RevId: 626074066
    maxtext authors committed Apr 18, 2024
    Configuration menu
    Copy the full SHA
    f52e6f7 View commit details
    Browse the repository at this point in the history

  3. Removes batch size from prefill attention calculation.

    @patemotter
    patemotter committed Apr 18, 2024
    Configuration menu
    Copy the full SHA
    f9498fa View commit details
    Browse the repository at this point in the history

  4. Fixes for inf testing.

    @patemotter
    patemotter committed Apr 18, 2024
    Configuration menu
    Copy the full SHA
    8a8e642 View commit details
    Browse the repository at this point in the history

  5. Revert "Fixes for inf testing." 

    This reverts commit b15b1d5.
    @patemotter
    patemotter committed Apr 18, 2024
    Configuration menu
    Copy the full SHA
    c2c3bce View commit details
    Browse the repository at this point in the history

  6. Fixes

    @patemotter
    patemotter committed Apr 18, 2024
    Configuration menu
    Copy the full SHA
    6c355d7 View commit details
    Browse the repository at this point in the history

Commits on Apr 19, 2024

  1. Fix subset of hosts dataloading

    @khatwanimohit
    khatwanimohit committed Apr 19, 2024
    Configuration menu
    Copy the full SHA
    20f2a0d View commit details
    Browse the repository at this point in the history

  2. Merge pull request #586 from google:mohit/subset_v4 

    PiperOrigin-RevId: 626387060
    maxtext authors committed Apr 19, 2024
    Configuration menu
    Copy the full SHA
    6ec7556 View commit details
    Browse the repository at this point in the history

  3. inference microbenchmark 

      - allow run specified stages
  - allow run specific prefill length(s)
  - delete prefill result
  - printout prefill result

added funcs in max_utils
    @morgandu
    morgandu committed Apr 19, 2024
    Configuration menu
    Copy the full SHA
    b46783c View commit details
    Browse the repository at this point in the history

  4. Merge branch 'main' into patemotter_attn_calc_fix

    @patemotter
    patemotter committed Apr 19, 2024
    Configuration menu
    Copy the full SHA
    2806017 View commit details
    Browse the repository at this point in the history

  5. Merge pull request #597 from google:mor--inference 

    PiperOrigin-RevId: 626410509
    maxtext authors committed Apr 19, 2024
    Configuration menu
    Copy the full SHA
    0e1c078 View commit details
    Browse the repository at this point in the history

  6. Merge pull request #600 from google:patemotter_attn_calc_fix 

    PiperOrigin-RevId: 626415915
    maxtext authors committed Apr 19, 2024
    Configuration menu
    Copy the full SHA
    1377756 View commit details
    Browse the repository at this point in the history

  7. Update Run_MaxText_via_xpk.md 

    Fixing typo.
    @RoshaniN
    RoshaniN committed Apr 19, 2024
    Configuration menu
    Copy the full SHA
    d3f488a View commit details
    Browse the repository at this point in the history

  8. Merge pull request #608 from google:RoshaniN-patch-1 

    PiperOrigin-RevId: 626446091
    maxtext authors committed Apr 19, 2024
    Configuration menu
    Copy the full SHA
    25adb3d View commit details
    Browse the repository at this point in the history

Commits on Apr 20, 2024

  1. inference_microbenchmark: 

      - time prefill only
  - benchmark prefill and insert
    @morgandu
    morgandu committed Apr 20, 2024
    Configuration menu
    Copy the full SHA
    c2e5b5e View commit details
    Browse the repository at this point in the history

Commits on Apr 22, 2024

  1. Mark nvidia devtools repo as trusted 

    This is a stopgaps measure to circumvent the nvidia repo's gpg signature issue
    @chajath
    chajath committed Apr 22, 2024
    Configuration menu
    Copy the full SHA
    5a3d5bd View commit details
    Browse the repository at this point in the history

  2. Merge pull request #610 from google:mor--inference 

    PiperOrigin-RevId: 627082389
    maxtext authors committed Apr 22, 2024
    Configuration menu
    Copy the full SHA
    c0bef1c View commit details
    Browse the repository at this point in the history

  3. Merge pull request #611 from google:yiinho-nvidia-devtools-trust 

    PiperOrigin-RevId: 627085448
    maxtext authors committed Apr 22, 2024
    Configuration menu
    Copy the full SHA
    49be620 View commit details
    Browse the repository at this point in the history

Commits on Apr 23, 2024

  1. Explicitly set AQT Freezer mode in MaxText. 

    PiperOrigin-RevId: 627250589
    @cdh4696
    cdh4696 authored and maxtext authors committed Apr 23, 2024
    Configuration menu
    Copy the full SHA
    a3d6d8f View commit details
    Browse the repository at this point in the history

  2. Move aqtp pin up

    @gobbleturk
    gobbleturk committed Apr 23, 2024
    Configuration menu
    Copy the full SHA
    3a14e47 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #615 from google:mattdavidow-move-aqtp-pin 

    PiperOrigin-RevId: 627417282
    maxtext authors committed Apr 23, 2024
    Configuration menu
    Copy the full SHA
    718d9e7 View commit details
    Browse the repository at this point in the history

  4. Pre-commit config

    @khatwanimohit
    khatwanimohit committed Apr 23, 2024
    Configuration menu
    Copy the full SHA
    ad072dd View commit details
    Browse the repository at this point in the history

Commits on Apr 24, 2024

  1. Update 128B config on v5e to use qkv_proj_offloaded remat_policy

    @raymondzouu
    raymondzouu committed Apr 24, 2024
    Configuration menu
    Copy the full SHA
    62b5012 View commit details
    Browse the repository at this point in the history

  2. Merge pull request #619 from google:raymondzou-128b-offload-fix 

    PiperOrigin-RevId: 627768373
    maxtext authors committed Apr 24, 2024
    Configuration menu
    Copy the full SHA
    206e84c View commit details
    Browse the repository at this point in the history

  3. Merge pull request #602 from google:mohit/codespell 

    PiperOrigin-RevId: 627781735
    maxtext authors committed Apr 24, 2024
    Configuration menu
    Copy the full SHA
    35f2b8a View commit details
    Browse the repository at this point in the history

  4. [MaxText] Rename llama2_7b_single_host_gpu.yml to make it clear that … 

    …it can be used for any number of host.

PiperOrigin-RevId: 627804089
    @bixia1
    bixia1 authored and maxtext authors committed Apr 24, 2024
    Configuration menu
    Copy the full SHA
    db31dd4 View commit details
    Browse the repository at this point in the history

Commits on Apr 25, 2024

  1. Split Mixtral test into two scripts

    @RissyRan
    RissyRan committed Apr 25, 2024
    Configuration menu
    Copy the full SHA
    548f639 View commit details
    Browse the repository at this point in the history

  2. Merge pull request #616 from RissyRan:separate_mixtral 

    PiperOrigin-RevId: 628127953
    maxtext authors committed Apr 25, 2024
    Configuration menu
    Copy the full SHA
    f6060b0 View commit details
    Browse the repository at this point in the history

  3. Update jax.tree_map to jax.tree_util.tree_map

    @RissyRan
    RissyRan committed Apr 25, 2024
    Configuration menu
    Copy the full SHA
    967941b View commit details
    Browse the repository at this point in the history

  4. Merge pull request #621 from RissyRan:tree_map 

    PiperOrigin-RevId: 628162121
    maxtext authors committed Apr 25, 2024
    Configuration menu
    Copy the full SHA
    18ba1a7 View commit details
    Browse the repository at this point in the history

Commits on Apr 26, 2024

  1. change norm sharding 

    fix lint

Revert "fix lint"

This reverts commit d8dc450.

fix lint
    @ZhiyuLi-goog
    ZhiyuLi-goog committed Apr 26, 2024
    Configuration menu
    Copy the full SHA
    9feab51 View commit details
    Browse the repository at this point in the history

  2. Merge pull request #623 from google:lizhiyu/change_norm_sharding 

    PiperOrigin-RevId: 628504084
    maxtext authors committed Apr 26, 2024
    Configuration menu
    Copy the full SHA
    6570445 View commit details
    Browse the repository at this point in the history

  3. Change l2norm to use jnp.sqrt

    @raymondzouu
    raymondzouu committed Apr 26, 2024
    Configuration menu
    Copy the full SHA
    50ba12e View commit details
    Browse the repository at this point in the history

Commits on Apr 29, 2024

  1. Merge pull request #618 from google:raymondzou-l2norm 

    PiperOrigin-RevId: 629072283
    maxtext authors committed Apr 29, 2024
    Configuration menu
    Copy the full SHA
    29e59de View commit details
    Browse the repository at this point in the history

Commits on Apr 30, 2024

  1. Fix test_tokenize

    @khatwanimohit
    khatwanimohit committed Apr 30, 2024
    Configuration menu
    Copy the full SHA
    4fc85d1 View commit details
    Browse the repository at this point in the history

  2. Streamlined setup.sh to have fewer apt install calls

    Rafi Witten committed Apr 30, 2024
    Configuration menu
    Copy the full SHA
    dab1ee1 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #629 from google:mohit/fix_token 

    PiperOrigin-RevId: 629530648
    maxtext authors committed Apr 30, 2024
    Configuration menu
    Copy the full SHA
    4062f51 View commit details
    Browse the repository at this point in the history

  4. loosen tolerance in assert_params_sufficiently_sharded

    @ZhiyuLi-goog
    ZhiyuLi-goog committed Apr 30, 2024
    Configuration menu
    Copy the full SHA
    75c4b03 View commit details
    Browse the repository at this point in the history

  5. Enable entropy on multihost CPUs.

    @RoshaniN
    RoshaniN committed Apr 30, 2024
    Configuration menu
    Copy the full SHA
    21177d4 View commit details
    Browse the repository at this point in the history

  6. Merge pull request #628 from google:lizhiyu/fix_params_sufficiently_s… 

    …harded

PiperOrigin-RevId: 629549121
    maxtext authors committed Apr 30, 2024
    Configuration menu
    Copy the full SHA
    b098e6a View commit details
    Browse the repository at this point in the history

  7. Merge pull request #631 from google:roshanin_entropy_2 

    PiperOrigin-RevId: 629553742
    maxtext authors committed Apr 30, 2024
    Configuration menu
    Copy the full SHA
    92f3abf View commit details
    Browse the repository at this point in the history

  8. Merge branch 'main' into rwitten_setup_streamlined

    Rafi Witten committed Apr 30, 2024
    Configuration menu
    Copy the full SHA
    d9ccd1e View commit details
    Browse the repository at this point in the history

  9. Merge pull request #633 from google:rwitten_setup_streamlined 

    PiperOrigin-RevId: 629566244
    maxtext authors committed Apr 30, 2024
    Configuration menu
    Copy the full SHA
    8f132d6 View commit details
    Browse the repository at this point in the history

Commits on May 1, 2024

  1. Add tests to GPU runner

    @michelle-yooh
    michelle-yooh committed May 1, 2024
    Configuration menu
    Copy the full SHA
    8a6f30d View commit details
    Browse the repository at this point in the history

  2. Replace deprecated np.product with np.prod

    @gobbleturk
    gobbleturk committed May 1, 2024
    Configuration menu
    Copy the full SHA
    3075bbe View commit details
    Browse the repository at this point in the history

  3. Merge pull request #630 from google:yooh/gpu-unit-tests 

    PiperOrigin-RevId: 629837420
    maxtext authors committed May 1, 2024
    Configuration menu
    Copy the full SHA
    fcf48fe View commit details
    Browse the repository at this point in the history

Commits on May 2, 2024

  1. fix norm sharding

    @ZhiyuLi-goog
    ZhiyuLi-goog committed May 2, 2024
    Configuration menu
    Copy the full SHA
    bc36642 View commit details
    Browse the repository at this point in the history

  2. Merge pull request #637 from google:lizhiyu/miss_sharding 

    PiperOrigin-RevId: 630121986
    maxtext authors committed May 2, 2024
    Configuration menu
    Copy the full SHA
    2ac0af9 View commit details
    Browse the repository at this point in the history

Commits on May 3, 2024

  1. Add Llama2-70b test

    @A9isha
    A9isha committed May 3, 2024
    Configuration menu
    Copy the full SHA
    982fe59 View commit details
    Browse the repository at this point in the history

  2. Merge pull request #639 from google:anisha-llama2-70b-test 

    PiperOrigin-RevId: 630417921
    maxtext authors committed May 3, 2024
    1 Configuration menu
    Copy the full SHA
    ef7cf7a View commit details
    Browse the repository at this point in the history

  3. Internal change only. 

    PiperOrigin-RevId: 630446330
    maxtext authors committed May 3, 2024
    1 Configuration menu
    Copy the full SHA
    e8b53e5 View commit details
    Browse the repository at this point in the history

  4. Add more tests for Mixtral

    @RissyRan
    RissyRan committed May 3, 2024
    Configuration menu
    Copy the full SHA
    ffcd34c View commit details
    Browse the repository at this point in the history

  5. Merge pull request #627 from google:ranran_more_moe 

    PiperOrigin-RevId: 630522953
    maxtext authors committed May 3, 2024
    1 Configuration menu
    Copy the full SHA
    bdb1b3e View commit details
    Browse the repository at this point in the history

Commits on May 6, 2024

  1. Make some AQT dataclasses to use keyword-only fields (1/N) 

    This cl introduces an temporary decorator that will be temporarily used during this migration. The eventual goal is to enforce kw_only=True in all dataclasses unless it's not feasible, aiming to make AQT less error-prune and improve readability.

PiperOrigin-RevId: 631132072
    maxtext authors committed May 6, 2024
    1 Configuration menu
    Copy the full SHA
    a519018 View commit details
    Browse the repository at this point in the history

Commits on May 7, 2024

  1. Reverts e8b53e5 

    PiperOrigin-RevId: 631465526
    @golechwierowicz
    golechwierowicz authored and maxtext authors committed May 7, 2024
    1 Configuration menu
    Copy the full SHA
    a28f518 View commit details
    Browse the repository at this point in the history

  2. Merge pull request #635 from google:mattdavidow-numpy-prod 

    PiperOrigin-RevId: 631470826
    maxtext authors committed May 7, 2024
    1 Configuration menu
    Copy the full SHA
    d590328 View commit details
    Browse the repository at this point in the history

Commits on May 9, 2024

  1. Update tflops calculation

    @RissyRan
    RissyRan committed May 9, 2024
    Configuration menu
    Copy the full SHA
    353fb83 View commit details
    Browse the repository at this point in the history

  2. fix sharding on generate cache in prefill results.

    @jwyang-google
    jwyang-google committed May 9, 2024
    Configuration menu
    Copy the full SHA
    f53203a View commit details
    Browse the repository at this point in the history

Commits on May 10, 2024

  1. Remove async XLA_FLAGS from A3 configs. 

    XLA PR openxla/xla#11422 removed some XLA flags relating to async collectives. This caused the A3 configs to fail to run, so this change removes such flags from the A3 configs. The flags removed are:

--xla_gpu_enable_async_all_gather=true
--xla_gpu_enable_async_reduce_scatter=true
--xla_gpu_enable_async_all_reduce=true

Such flags had no impact before the XLA PR as the async collectives were already enabled by default.
    @reedwm
    reedwm committed May 10, 2024
    Configuration menu
    Copy the full SHA
    e84114d View commit details
    Browse the repository at this point in the history

  2. Update llama2_7b_gpu.yml 

    PiperOrigin-RevId: 631752008
    maxtext authors committed May 10, 2024
    1 Configuration menu
    Copy the full SHA
    ef00f3b View commit details
    Browse the repository at this point in the history

  3. Merge pull request #644 from google:moe_param 

    PiperOrigin-RevId: 632315686
    maxtext authors committed May 10, 2024
    1 Configuration menu
    Copy the full SHA
    f9c9dd8 View commit details
    Browse the repository at this point in the history

  4. Merge pull request #641 from google:prefill-optimize 

    PiperOrigin-RevId: 632517062
    maxtext authors committed May 10, 2024
    1 Configuration menu
    Copy the full SHA
    e327805 View commit details
    Browse the repository at this point in the history

  5. Add forward pass logit check test for Llama2-7b

    @A9isha
    A9isha committed May 10, 2024
    Configuration menu
    Copy the full SHA
    7434359 View commit details
    Browse the repository at this point in the history

  6. Merge pull request #643 from google:anisha-debug-decoding 

    PiperOrigin-RevId: 632613803
    maxtext authors committed May 10, 2024
    1 Configuration menu
    Copy the full SHA
    4f3a0d3 View commit details
    Browse the repository at this point in the history

Commits on May 13, 2024

  1. Eval the command string from XPK for GPU script

    @michelle-yooh
    michelle-yooh committed May 13, 2024
    Configuration menu
    Copy the full SHA
    4f526a8 View commit details
    Browse the repository at this point in the history

Commits on May 14, 2024

  1. Merge pull request #640 from google:yooh/add_oss_tests 

    PiperOrigin-RevId: 633628393
    maxtext authors committed May 14, 2024
    1 Configuration menu
    Copy the full SHA
    87c6430 View commit details
    Browse the repository at this point in the history

  2. Remove cases where the deprecated --xla_gpu_simplify_all_fp_conversio… 

    …ns is set to its default value.

PiperOrigin-RevId: 633645462
    @dimitar-asenov
    dimitar-asenov authored and maxtext authors committed May 14, 2024
    1 Configuration menu
    Copy the full SHA
    618450b View commit details
    Browse the repository at this point in the history

Commits on May 15, 2024

  1. streamline CI test structure

    @michelle-yooh
    michelle-yooh committed May 15, 2024
    Configuration menu
    Copy the full SHA
    2347382 View commit details
    Browse the repository at this point in the history

  2. Merge branch 'main' into xla_flags_fix

    @reedwm
    reedwm committed May 15, 2024
    Configuration menu
    Copy the full SHA
    2d2c3c4 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #645 from reedwm:xla_flags_fix 

    PiperOrigin-RevId: 634085946
    maxtext authors committed May 15, 2024
    Configuration menu
    Copy the full SHA
    d9aaaab View commit details
    Browse the repository at this point in the history

  4. fix pylint 

    fix pylint: Using variable 'p_eval_step' before assignment (#651)
    @aireenmei
    aireenmei committed May 15, 2024
    Configuration menu
    Copy the full SHA
    0a626ed View commit details
    Browse the repository at this point in the history

  5. Merge pull request #642 from google:yooh/streamline-unit-tests 

    PiperOrigin-RevId: 634120234
    maxtext authors committed May 15, 2024
    Configuration menu
    Copy the full SHA
    cd48a14 View commit details
    Browse the repository at this point in the history

  6. Merge pull request #649 from google:aireen/fix_pylint 

    PiperOrigin-RevId: 634121112
    maxtext authors committed May 15, 2024
    Configuration menu
    Copy the full SHA
    3e65eb6 View commit details
    Browse the repository at this point in the history

Commits on May 16, 2024

  1. Remove async XLA_FLAGS from A3 configs

    @michelle-yooh
    michelle-yooh committed May 16, 2024
    Configuration menu
    Copy the full SHA
    b98ad6c View commit details
    Browse the repository at this point in the history

  2. Add llama-70b gpu config. 

    PiperOrigin-RevId: 634267313
    @golechwierowicz
    golechwierowicz authored and maxtext authors committed May 16, 2024
    Configuration menu
    Copy the full SHA
    fa92a21 View commit details
    Browse the repository at this point in the history

Commits on May 17, 2024

  1. Merge pull request #652 from google:yooh/xla-flags-fix 

    PiperOrigin-RevId: 634572059
    maxtext authors committed May 17, 2024
    Configuration menu
    Copy the full SHA
    8930215 View commit details
    Browse the repository at this point in the history

Commits on May 20, 2024

  1. Support data input from HuggingFace

    @aireenmei
    aireenmei committed May 20, 2024
    Configuration menu
    Copy the full SHA
    3face14 View commit details
    Browse the repository at this point in the history

  2. Update the NCCL flags for A3+.

    @yangyuwei
    yangyuwei committed May 20, 2024
    Configuration menu
    Copy the full SHA
    65fbd7d View commit details
    Browse the repository at this point in the history

  3. Merge pull request #650 from google:yangyuwei-maxtext 

    PiperOrigin-RevId: 635517779
    maxtext authors committed May 20, 2024
    Configuration menu
    Copy the full SHA
    e60cabf View commit details
    Browse the repository at this point in the history

  4. add gemma logit test

    @A9isha
    A9isha committed May 20, 2024
    Configuration menu
    Copy the full SHA
    cc063cf View commit details
    Browse the repository at this point in the history

  5. Merge pull request #654 from google:anisha-gemma2b-fwd 

    PiperOrigin-RevId: 635527190
    maxtext authors committed May 20, 2024
    Configuration menu
    Copy the full SHA
    3519271 View commit details
    Browse the repository at this point in the history

Commits on May 21, 2024

  1. Integrate orbax logger in maxtext for structured logging. 

    Integrate orbax logger in maxtext for structured logging.

Integrate orbax logger in maxtext for structured logging.

Integrate orbax logger in maxtext for structured logging.

Integrate orbax logger in maxtext for structured logging.

Integrate orbax logger in maxtext for structured logging.
    @abhinavclemson
    abhinavclemson committed May 21, 2024
    Configuration menu
    Copy the full SHA
    1f34741 View commit details
    Browse the repository at this point in the history

  2. Merge pull request #646 from google:aireen/hf_data_pr2 

    PiperOrigin-RevId: 635853912
    maxtext authors committed May 21, 2024
    Configuration menu
    Copy the full SHA
    6669c00 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #658 from google:orbax-logger-abhinav 

    PiperOrigin-RevId: 635909655
    maxtext authors committed May 21, 2024
    Configuration menu
    Copy the full SHA
    6542d25 View commit details
    Browse the repository at this point in the history

Commits on May 22, 2024

  1. fix hf input pipeline

    @aireenmei
    aireenmei committed May 22, 2024
    Configuration menu
    Copy the full SHA
    76be151 View commit details
    Browse the repository at this point in the history

  2. Merge pull request #664 from google:aireen/fix_hf 

    PiperOrigin-RevId: 636035816
    maxtext authors committed May 22, 2024
    Configuration menu
    Copy the full SHA
    15f11c5 View commit details
    Browse the repository at this point in the history

  3. Fix prefill assertion

    @RissyRan
    RissyRan committed May 22, 2024
    Configuration menu
    Copy the full SHA
    91c73e4 View commit details
    Browse the repository at this point in the history

  4. Merge pull request #661 from google:fix_prefill 

    PiperOrigin-RevId: 636279640
    maxtext authors committed May 22, 2024
    Configuration menu
    Copy the full SHA
    b70ec2a View commit details
    Browse the repository at this point in the history

Commits on May 23, 2024

  1. Remove decode asserts from Gemma test files

    @khatwanimohit
    khatwanimohit committed May 23, 2024
    Configuration menu
    Copy the full SHA
    ed15011 View commit details
    Browse the repository at this point in the history

  2. Merge pull request #668 from google:mohit/fix_gemma2b 

    PiperOrigin-RevId: 636566491
    maxtext authors committed May 23, 2024
    Configuration menu
    Copy the full SHA
    af36df3 View commit details
    Browse the repository at this point in the history

Commits on May 24, 2024

  1. add single controller flag

    @sadikneipp
    sadikneipp committed May 24, 2024
    Configuration menu
    Copy the full SHA
    3da83da View commit details
    Browse the repository at this point in the history

  2. fix OOM issue running inference microbenchmark with llama13b on v5e4

    @jwyang-google
    jwyang-google committed May 24, 2024
    Configuration menu
    Copy the full SHA
    e2e91f9 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #669 from google:ksadi/patch-1 

    PiperOrigin-RevId: 637028379
    maxtext authors committed May 24, 2024
    Configuration menu
    Copy the full SHA
    6168580 View commit details
    Browse the repository at this point in the history

Commits on May 25, 2024

  1. Merge branch 'main' into inference-microbenchmark-fix

    @jwyang-google
    jwyang-google committed May 25, 2024
    Configuration menu
    Copy the full SHA
    e36ea81 View commit details
    Browse the repository at this point in the history

Commits on May 26, 2024

  1. Merge pull request #670 from google:inference-microbenchmark-fix 

    PiperOrigin-RevId: 637329632
    maxtext authors committed May 26, 2024
    Configuration menu
    Copy the full SHA
    ac7abce View commit details
    Browse the repository at this point in the history

Commits on May 28, 2024

  1. Add Llama2 13B Tests

    @morgandu
    morgandu committed May 28, 2024
    Configuration menu
    Copy the full SHA
    c2fb0fd View commit details
    Browse the repository at this point in the history

  2. Merge pull request #565 from google:mor--llama2-13b-test 

    PiperOrigin-RevId: 637997076
    maxtext authors committed May 28, 2024
    Configuration menu
    Copy the full SHA
    313d31b View commit details
    Browse the repository at this point in the history

Commits on May 29, 2024

  1. Don't clip fp8 stats

    @anfals
    anfals committed May 29, 2024
    Configuration menu
    Copy the full SHA
    b061356 View commit details
    Browse the repository at this point in the history

  2. Merge pull request #672 from google:anfals/fp8_clipping_fix 

    PiperOrigin-RevId: 638321669
    maxtext authors committed May 29, 2024
    Configuration menu
    Copy the full SHA
    c47a656 View commit details
    Browse the repository at this point in the history

  3. Integrate nsys profiler 

    Remove 'enable_profiler' config and add 'profiler' config instead
    @michelle-yooh
    michelle-yooh committed May 29, 2024
    Configuration menu
    Copy the full SHA
    d05052e View commit details
    Browse the repository at this point in the history

  4. Merge pull request #663 from google:yooh/integrate-nsys 

    PiperOrigin-RevId: 638403056
    maxtext authors committed May 29, 2024
    Configuration menu
    Copy the full SHA
    48f2524 View commit details
    Browse the repository at this point in the history

Commits on May 30, 2024

  1. Add MoE matmul implementation

    @RissyRan
    RissyRan committed May 30, 2024
    Configuration menu
    Copy the full SHA
    851d048 View commit details
    Browse the repository at this point in the history

Commits on May 31, 2024

  1. Merge pull request #659 from google:ranran_moe 

    PiperOrigin-RevId: 639085670
    maxtext authors committed May 31, 2024
    Configuration menu
    Copy the full SHA
    410d237 View commit details
    Browse the repository at this point in the history

  2. fix OUTPUT_PATH in v5e/128b.sh

    @ZhiyuLi-goog
    ZhiyuLi-goog committed May 31, 2024
    Configuration menu
    Copy the full SHA
    3687526 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #676 from google:lizhiyu/fix_path 

    PiperOrigin-RevId: 639128398
    maxtext authors committed May 31, 2024
    Configuration menu
    Copy the full SHA
    52b954b View commit details
    Browse the repository at this point in the history

  4. squash

    @Bslabe123
    Bslabe123 committed May 31, 2024
    Configuration menu
    Copy the full SHA
    ef96d6d View commit details
    Browse the repository at this point in the history

  5. Merge pull request #665 from google:prometheus-port-flag 

    PiperOrigin-RevId: 639145146
    maxtext authors committed May 31, 2024
    Configuration menu
    Copy the full SHA
    3a441e8 View commit details
    Browse the repository at this point in the history

  6. Update flops calculation to active experts in moe

    @RissyRan
    RissyRan committed May 31, 2024
    Configuration menu
    Copy the full SHA
    288912b View commit details
    Browse the repository at this point in the history

  7. Merge pull request #678 from google:tflop_moe 

    PiperOrigin-RevId: 639194660
    maxtext authors committed May 31, 2024
    Configuration menu
    Copy the full SHA
    f12ba54 View commit details
    Browse the repository at this point in the history

Commits on Jun 3, 2024

  1. Enable kv cache layout control

    @morgandu
    morgandu committed Jun 3, 2024
    Configuration menu
    Copy the full SHA
    57429da View commit details
    Browse the repository at this point in the history

  2. Merge pull request #667 from google:mor--kv-cache-layout 

    PiperOrigin-RevId: 639790336
    maxtext authors committed Jun 3, 2024
    Configuration menu
    Copy the full SHA
    34412d4 View commit details
    Browse the repository at this point in the history

  3. Fix Gemma Readme link

    @khatwanimohit
    khatwanimohit committed Jun 3, 2024
    Configuration menu
    Copy the full SHA
    d6b721c View commit details
    Browse the repository at this point in the history

  4. Internal change only. 

    Reverts a28f518

PiperOrigin-RevId: 639890999
    maxtext authors committed Jun 3, 2024
    Configuration menu
    Copy the full SHA
    ca55917 View commit details
    Browse the repository at this point in the history

  5. Upgrade Pinned Base Image for GPU

    @anfals
    anfals committed Jun 3, 2024
    Configuration menu
    Copy the full SHA
    5723009 View commit details
    Browse the repository at this point in the history

  6. Metrics bug: server_lib should be config_lib

    @Bslabe123
    Bslabe123 committed Jun 3, 2024
    Configuration menu
    Copy the full SHA
    83159ab View commit details
    Browse the repository at this point in the history

Commits on Jun 4, 2024

  1. Merge pull request #681 from google:prometheus-flag-import-fix 

    PiperOrigin-RevId: 640138468
    maxtext authors committed Jun 4, 2024
    Configuration menu
    Copy the full SHA
    b7f5864 View commit details
    Browse the repository at this point in the history

  2. Merge pull request #682 from google:anfals/up_gpu_base_image 

    PiperOrigin-RevId: 640180246
    maxtext authors committed Jun 4, 2024
    Configuration menu
    Copy the full SHA
    a77e867 View commit details
    Browse the repository at this point in the history

Commits on Jun 5, 2024

  1. Fix MoE matmul scale issue

    @RissyRan
    RissyRan committed Jun 5, 2024
    Configuration menu
    Copy the full SHA
    cdb4853 View commit details
    Browse the repository at this point in the history

  2. Removed unused Pallas import from layers/attentions.py 

    PiperOrigin-RevId: 640481280
    @superbobry
    superbobry authored and maxtext authors committed Jun 5, 2024
    Configuration menu
    Copy the full SHA
    2a6154f View commit details
    Browse the repository at this point in the history

  3. Change norm sharding for llama2-7b to fsdp. 

    PiperOrigin-RevId: 640498890
    @golechwierowicz
    golechwierowicz authored and maxtext authors committed Jun 5, 2024
    Configuration menu
    Copy the full SHA
    cb2c69b View commit details
    Browse the repository at this point in the history

  4. Copybara import of the project: 

    --
d7d694f by RissyRan <[email protected]>:

Fix forward test for Mixtral

COPYBARA_INTEGRATE_REVIEW=#679 from google:ranran_fix_forward_test d7d694f
PiperOrigin-RevId: 640537456
    @RissyRan
    RissyRan authored and maxtext authors committed Jun 5, 2024
    Configuration menu
    Copy the full SHA
    d185209 View commit details
    Browse the repository at this point in the history

  5. Set additional flags for a3 and a3plus

    @michelle-yooh
    michelle-yooh committed Jun 5, 2024
    Configuration menu
    Copy the full SHA
    eafc6af View commit details
    Browse the repository at this point in the history

  6. Use run_id instead of sha for docker tag

    @jonb377
    jonb377 committed Jun 5, 2024
    Configuration menu
    Copy the full SHA
    c76e92b View commit details
    Browse the repository at this point in the history

  7. Merge pull request #685 from google:fix_matmul_scale 

    PiperOrigin-RevId: 640594189
    maxtext authors committed Jun 5, 2024
    Configuration menu
    Copy the full SHA
    16aeac4 View commit details
    Browse the repository at this point in the history

  8. Merge branch 'main' into jonbolin/gha

    @jonb377
    jonb377 committed Jun 5, 2024
    Configuration menu
    Copy the full SHA
    a4fddec View commit details
    Browse the repository at this point in the history

  9. Merge pull request #684 from google:yooh/a3plus-flags 

    PiperOrigin-RevId: 640650131
    maxtext authors committed Jun 5, 2024
    Configuration menu
    Copy the full SHA
    c0ad118 View commit details
    Browse the repository at this point in the history

  10. Merge branch 'main' into jonbolin/gha

    @jonb377
    jonb377 committed Jun 5, 2024
    Configuration menu
    Copy the full SHA
    24836aa View commit details
    Browse the repository at this point in the history

Commits on Jun 6, 2024

  1. refactor data input pipeline and add perf data

    @aireenmei
    aireenmei committed Jun 6, 2024
    Configuration menu
    Copy the full SHA
    d6fae01 View commit details
    Browse the repository at this point in the history

Commits on Jun 7, 2024

  1. Merge pull request #688 from google:jonbolin/gha 

    PiperOrigin-RevId: 641367903
    maxtext authors committed Jun 7, 2024
    Configuration menu
    Copy the full SHA
    3e34b16 View commit details
    Browse the repository at this point in the history

  2. Add gpt3 175b on v5e config

    @raymondzouu
    raymondzouu committed Jun 7, 2024
    Configuration menu
    Copy the full SHA
    a8dd031 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #687 from google:raymondzou-gpt3-175b-config 

    PiperOrigin-RevId: 641391694
    maxtext authors committed Jun 7, 2024
    Configuration menu
    Copy the full SHA
    3ba16a9 View commit details
    Browse the repository at this point in the history

Commits on Jun 11, 2024

  1. Merge pull request #680 from google:aireen/grain_improve 

    PiperOrigin-RevId: 642384555
    maxtext authors committed Jun 11, 2024
    Configuration menu
    Copy the full SHA
    316dbca View commit details
    Browse the repository at this point in the history

Commits on Jun 12, 2024

  1. Pipeline parallelism support (linear only)

    @gobbleturk
    gobbleturk committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    5ae05cc View commit details
    Browse the repository at this point in the history

  2. Turn on layer scanning for llama2-7b on GPU. 

    This better utilizes recent optimizations to collective approximation in the XLA latency hiding scheduler.

PiperOrigin-RevId: 642559284
    @golechwierowicz
    golechwierowicz authored and maxtext authors committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    a2fdf29 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #691 from google:mattdavidow-pipeline-linear 

    PiperOrigin-RevId: 642648106
    maxtext authors committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    7cdca96 View commit details
    Browse the repository at this point in the history

  4. reshape q

    @morgandu
    morgandu committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    02b681e View commit details
    Browse the repository at this point in the history

  5. Merge pull request #690 from google:mor--reshape-q 

    PiperOrigin-RevId: 642676990
    maxtext authors committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    d0701b0 View commit details
    Browse the repository at this point in the history

  6. Add profiler flags to JetStream server 

    Add jetstream config

backward compatible
    @JoeZijunZhou
    JoeZijunZhou committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    e682651 View commit details
    Browse the repository at this point in the history

  7. fix tfds instruction

    @aireenmei
    aireenmei committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    3d64c10 View commit details
    Browse the repository at this point in the history

  8. Add vanilla megablox to MoE

    @RissyRan
    RissyRan committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    0fc492d View commit details
    Browse the repository at this point in the history

Commits on Jun 13, 2024

  1. Merge pull request #689 from google:vanilla_megablox 

    PiperOrigin-RevId: 642798191
    maxtext authors committed Jun 13, 2024
    Configuration menu
    Copy the full SHA
    e898606 View commit details
    Browse the repository at this point in the history

  2. Add llama2 70b training config for v5e

    @raymondzouu
    raymondzouu committed Jun 13, 2024
    Configuration menu
    Copy the full SHA
    f5ede57 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #692 from google:zijun/inference-profiler 

    PiperOrigin-RevId: 643127965
    maxtext authors committed Jun 13, 2024
    Configuration menu
    Copy the full SHA
    8c11811 View commit details
    Browse the repository at this point in the history

  4. Merge pull request #695 from google:raymondzou-llama2-70b-config 

    PiperOrigin-RevId: 643131869
    maxtext authors committed Jun 13, 2024
    Configuration menu
    Copy the full SHA
    d6e332c View commit details
    Browse the repository at this point in the history

  5. Merge pull request #702 from google:aireen/doc_fix 

    PiperOrigin-RevId: 643152508
    maxtext authors committed Jun 13, 2024
    Configuration menu
    Copy the full SHA
    f50f5ed View commit details
    Browse the repository at this point in the history

Commits on Jun 17, 2024

  1. base.yml changes 

    circular changes to pipeline.py

pyconfig circ changes

pipeline parallel tests circular style

tree map, half passed tests

Total iterations circularized

improved iteration comment

run all tests

test both circular and non-circular

circ storage comment

circ storage pushing index comment
    @gobbleturk
    gobbleturk committed Jun 17, 2024
    Configuration menu
    Copy the full SHA
    a75d9a9 View commit details
    Browse the repository at this point in the history

  2. Account for new mesh axes for llama2-7b, and llama2-70b on GPUs. 

    PiperOrigin-RevId: 643999933
    @golechwierowicz
    golechwierowicz authored and maxtext authors committed Jun 17, 2024
    Configuration menu
    Copy the full SHA
    5c9e569 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #701 from google:mattdavidow-pipeline-circular 

    PiperOrigin-RevId: 644065470
    maxtext authors committed Jun 17, 2024
    Configuration menu
    Copy the full SHA
    75b3a5e View commit details
    Browse the repository at this point in the history

Commits on Jun 18, 2024

  1. Sharding the llama2 70b on v5e-16 more efficiently. 

    https://arxiv.org/pdf/2211.05102
https://arxiv.org/pdf/1909.08053
    Zhihao Shan committed Jun 18, 2024
    Configuration menu
    Copy the full SHA
    8bf9f8e View commit details
    Browse the repository at this point in the history

  2. Merge pull request #706 from google:zhihaoshan_dev 

    PiperOrigin-RevId: 644220882
    maxtext authors committed Jun 18, 2024
    Configuration menu
    Copy the full SHA
    180a780 View commit details
    Browse the repository at this point in the history

  3. add compute_axis_order

    @morgandu
    morgandu committed Jun 18, 2024
    Configuration menu
    Copy the full SHA
    5cf2e57 View commit details
    Browse the repository at this point in the history

  4. Merge pull request #709 from google:mor--compute-axis-order 

    PiperOrigin-RevId: 644386269
    maxtext authors committed Jun 18, 2024
    Configuration menu
    Copy the full SHA
    fe4bfdc View commit details
    Browse the repository at this point in the history

  5. Add maxengine_server configs to base.yml

    @gobbleturk
    gobbleturk committed Jun 18, 2024
    Configuration menu
    Copy the full SHA
    8a22374 View commit details
    Browse the repository at this point in the history

  6. Merge pull request #713 from google:mattdavidow-add-max-engine-flags-… 

    …to-base

PiperOrigin-RevId: 644472824
    maxtext authors committed Jun 18, 2024
    Configuration menu
    Copy the full SHA
    2d86110 View commit details
    Browse the repository at this point in the history

  7. Add FSDP + Megablox

    @RissyRan
    RissyRan committed Jun 18, 2024
    Configuration menu
    Copy the full SHA
    c7cf774 View commit details
    Browse the repository at this point in the history

Commits on Jun 19, 2024

  1. Merge pull request #705 from google:fsdp_megablox 

    PiperOrigin-RevId: 644577095
    maxtext authors committed Jun 19, 2024
    Configuration menu
    Copy the full SHA
    5d7b037 View commit details
    Browse the repository at this point in the history

Commits on Jun 20, 2024

  1. Llama3-8b model config

    @khatwanimohit
    khatwanimohit committed Jun 20, 2024
    Configuration menu
    Copy the full SHA
    6e90b9e View commit details
    Browse the repository at this point in the history

  2. Merge pull request #653 from google:mohit/llama3 

    PiperOrigin-RevId: 644954456
    maxtext authors committed Jun 20, 2024
    Configuration menu
    Copy the full SHA
    f7ee8c6 View commit details
    Browse the repository at this point in the history

  3. MaxText package

    @khatwanimohit
    khatwanimohit committed Jun 20, 2024
    Configuration menu
    Copy the full SHA
    e44e07a View commit details
    Browse the repository at this point in the history

  4. Merge pull request #707 from google:mohit/package 

    PiperOrigin-RevId: 645098382
    maxtext authors committed Jun 20, 2024
    Configuration menu
    Copy the full SHA
    d82a44f View commit details
    Browse the repository at this point in the history

Commits on Jun 21, 2024

  1. fix data loading from HF hub

    @aireenmei
    aireenmei committed Jun 21, 2024
    Configuration menu
    Copy the full SHA
    ea4d32a View commit details
    Browse the repository at this point in the history

  2. Fix llama2-{7,70}b sharding on GPU. 

    PiperOrigin-RevId: 645365795
    @golechwierowicz
    golechwierowicz authored and maxtext authors committed Jun 21, 2024
    Configuration menu
    Copy the full SHA
    0ea5f38 View commit details
    Browse the repository at this point in the history

  3. Move stage to second axis in mesh 

    Move stage to second axis in mesh
    @gobbleturk
    gobbleturk committed Jun 21, 2024
    Configuration menu
    Copy the full SHA
    8f49ace View commit details
    Browse the repository at this point in the history

  4. Merge pull request #717 from google:mattdavidow-stage-second 

    PiperOrigin-RevId: 645457324
    maxtext authors committed Jun 21, 2024
    Configuration menu
    Copy the full SHA
    592f45e View commit details
    Browse the repository at this point in the history

  5. Merge pull request #716 from google:aireen/fix_hf_hub 

    PiperOrigin-RevId: 645516055
    maxtext authors committed Jun 21, 2024
    Configuration menu
    Copy the full SHA
    9df5761 View commit details
    Browse the repository at this point in the history

Commits on Jun 22, 2024

  1. Copybara import of the project: 

    --
1718b89 by RissyRan <[email protected]>:

Refactor permute and unpermute operations

COPYBARA_INTEGRATE_REVIEW=#714 from google:refactor_mega b101cbc
PiperOrigin-RevId: 645591567
    @RissyRan
    RissyRan authored and maxtext authors committed Jun 22, 2024
    Configuration menu
    Copy the full SHA
    7a40096 View commit details
    Browse the repository at this point in the history

Commits on Jun 23, 2024

  1. Fix Mesh setup for multiprocess CPUs.

    @RoshaniN
    RoshaniN committed Jun 23, 2024
    Configuration menu
    Copy the full SHA
    24dc66c View commit details
    Browse the repository at this point in the history

Commits on Jun 24, 2024

  1. Merge pull request #723 from google:fix_jax_coordinator 

    PiperOrigin-RevId: 646137860
    maxtext authors committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    5db6d73 View commit details
    Browse the repository at this point in the history

  2. add kv_quant_axis

    @morgandu
    morgandu committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    8c79a64 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #708 from google:mor--compute-axis-order-n-quantiz… 

    …e-kv-over-hd

PiperOrigin-RevId: 646192065
    maxtext authors committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    a76f9d4 View commit details
    Browse the repository at this point in the history

  4. Add a directory check for the . If it fails, attempt to check a path … 

    …relative to the base config, similar to what is done for model configurations.

Minor update

Remove the raised exception
    Adam O'Brien committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    41f9af1 View commit details
    Browse the repository at this point in the history

  5. Add mistral tokenizer to maxtext/assets

    @vipannalla
    vipannalla committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    465f67a View commit details
    Browse the repository at this point in the history

  6. Merge pull request #715 from google:mxla-updates 

    PiperOrigin-RevId: 646238794
    maxtext authors committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    7a872f9 View commit details
    Browse the repository at this point in the history

  7. Merge branch 'main' into mistral_tok

    @vipannalla
    vipannalla committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    580d922 View commit details
    Browse the repository at this point in the history

  8. Merge pull request #719 from google:mistral_tok 

    PiperOrigin-RevId: 646255569
    maxtext authors committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    6b7d88c View commit details
    Browse the repository at this point in the history

  9. Update the dependencies to prepare for integration of emergency check… 

    …pointing

Withhold some package versions

Update version of typing_extensions
    @xuefgu
    xuefgu committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    91af148 View commit details
    Browse the repository at this point in the history

Commits on Jun 25, 2024

  1. Merge pull request #724 from google:requirements 

    PiperOrigin-RevId: 646312672
    maxtext authors committed Jun 25, 2024
    Configuration menu
    Copy the full SHA
    e7c1f01 View commit details
    Browse the repository at this point in the history

  2. Make broadcasting from one replica to all more memory efficient 

    PiperOrigin-RevId: 646526020
    maxtext authors committed Jun 25, 2024
    Configuration menu
    Copy the full SHA
    9482bf1 View commit details
    Browse the repository at this point in the history

  3. Inference Microbenchmark Sweep

    @yeandy
    yeandy committed Jun 25, 2024
    Configuration menu
    Copy the full SHA
    9606e62 View commit details
    Browse the repository at this point in the history

  4. Merge pull request #662 from google:mor--kv-cache-layout-reformat-output 

    PiperOrigin-RevId: 646637997
    maxtext authors committed Jun 25, 2024
    Configuration menu
    Copy the full SHA
    5a215db View commit details
    Browse the repository at this point in the history

Commits on Jun 26, 2024

  1. Fix mesh_axes and data_sharding for LLaMA 2 GPU configs. 

    PiperOrigin-RevId: 646795068
    @golechwierowicz
    golechwierowicz authored and maxtext authors committed Jun 26, 2024
    Configuration menu
    Copy the full SHA
    679ec8c View commit details
    Browse the repository at this point in the history

Commits on Jun 27, 2024

  1. Allow owners to have any approver 

    Fix AddLabel syntax

Fix punctuation
    @gobbleturk
    gobbleturk committed Jun 27, 2024
    Configuration menu
    Copy the full SHA
    ff2061f View commit details
    Browse the repository at this point in the history

  2. Enable saving using Orbax's emergency checkpoint manager 

    fix data loading from HF hub

Add explanation to the emergency checkpoint feature

Fix pylint issues

Minor changes to the config file

resolve conflicts

Inference Microbenchmark Sweep

Fix mesh_axes and data_sharding for LLaMA 2 GPU configs.

PiperOrigin-RevId: 646795068
    @xuefgu
    xuefgu committed Jun 27, 2024
    Configuration menu
    Copy the full SHA
    d4e6f25 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #722 from google:mattdavidow-fix-owners 

    PiperOrigin-RevId: 647360753
    maxtext authors committed Jun 27, 2024
    Configuration menu
    Copy the full SHA
    16903de View commit details
    Browse the repository at this point in the history

  4. Merge pull request #720 from google:emergency-checkpoint 

    PiperOrigin-RevId: 647390634
    maxtext authors committed Jun 27, 2024
    Configuration menu
    Copy the full SHA
    fe0dee1 View commit details
    Browse the repository at this point in the history

  5. Add Llama2 7B, 13B high performance training configs

    @raymondzouu
    raymondzouu committed Jun 27, 2024
    Configuration menu
    Copy the full SHA
    b5160d8 View commit details
    Browse the repository at this point in the history

  6. Load/Save Aqt quantized checkpoint.

    @singh-mitali
    singh-mitali committed Jun 27, 2024
    Configuration menu
    Copy the full SHA
    965c8af View commit details
    Browse the repository at this point in the history

  7. Merge pull request #726 from google:raymondzou-llama2 

    PiperOrigin-RevId: 647436094
    maxtext authors committed Jun 27, 2024
    Configuration menu
    Copy the full SHA
    8dedf9f View commit details
    Browse the repository at this point in the history

  8. Merge pull request #731 from google:msingh-aqt-ckpt 

    PiperOrigin-RevId: 647452066
    maxtext authors committed Jun 27, 2024
    Configuration menu
    Copy the full SHA
    e926938 View commit details
    Browse the repository at this point in the history

  9. modify prefill to return first token

    @jwyang-google
    jwyang-google committed Jun 27, 2024
    Configuration menu
    Copy the full SHA
    8900fe3 View commit details
    Browse the repository at this point in the history

Commits on Jun 28, 2024

  1. Merge pull request #727 from google:prefill-return-first-token 

    PiperOrigin-RevId: 647693612
    maxtext authors committed Jun 28, 2024
    Configuration menu
    Copy the full SHA
    d9138b1 View commit details
    Browse the repository at this point in the history

  2. Fix and protect simple_layer 

    Fix and protect simple_layer

Fix and protect simple_layer

Fix and protect simple_layer
    @gobbleturk
    gobbleturk committed Jun 28, 2024
    Configuration menu
    Copy the full SHA
    1128ed5 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #732 from google:mattdavidow-support-simple 

    PiperOrigin-RevId: 647752765
    maxtext authors committed Jun 28, 2024
    Configuration menu
    Copy the full SHA
    93efadf View commit details
    Browse the repository at this point in the history

Commits on Jun 29, 2024

  1. Adding option for int4 quantization to kvcache.

    @singh-mitali
    singh-mitali committed Jun 29, 2024
    Configuration menu
    Copy the full SHA
    4378403 View commit details
    Browse the repository at this point in the history

Commits on Jul 1, 2024

  1. support eval dataset and refactor

    @aireenmei
    aireenmei committed Jul 1, 2024
    Configuration menu
    Copy the full SHA
    77f079f View commit details
    Browse the repository at this point in the history

  2. Merge pull request #737 from google:msingh-kv 

    PiperOrigin-RevId: 648407356
    maxtext authors committed Jul 1, 2024
    Configuration menu
    Copy the full SHA
    bdeab2b View commit details
    Browse the repository at this point in the history

  3. Support partial overrides for logical_axis_rules.

    @golechwierowicz
    golechwierowicz committed Jul 1, 2024
    Configuration menu
    Copy the full SHA
    818cd06 View commit details
    Browse the repository at this point in the history

Commits on Jul 2, 2024

  1. Merge pull request #738 from google:aireen/tfds-eval 

    PiperOrigin-RevId: 648687112
    maxtext authors committed Jul 2, 2024
    Configuration menu
    Copy the full SHA
    f8ae413 View commit details
    Browse the repository at this point in the history

  2. Fix simple test step count

    @gobbleturk
    gobbleturk committed Jul 2, 2024
    Configuration menu
    Copy the full SHA
    ada70ff View commit details
    Browse the repository at this point in the history

  3. Merge pull request #745 from google:mattdavidow-fix-simple-test 

    PiperOrigin-RevId: 648835553
    maxtext authors committed Jul 2, 2024
    Configuration menu
    Copy the full SHA
    0868add View commit details
    Browse the repository at this point in the history

  4. Clean up MoE brute force implementation

    @RissyRan
    RissyRan committed Jul 2, 2024
    Configuration menu
    Copy the full SHA
    68b72b9 View commit details
    Browse the repository at this point in the history

  5. Preliminary restore with lots of hardcoding and hacking 

    Refactor the code and remove the hardcoding

More refactoring

Cleanup for pull request

Address linting issues

Preliminary restore with lots of hardcoding and hacking

Refactor the code and remove the hardcoding

More refactoring

Cleanup for pull request

Address linting issues

Small formatting

Fix merging issues
    @xuefgu
    xuefgu committed Jul 2, 2024
    Configuration menu
    Copy the full SHA
    e7c019a View commit details
    Browse the repository at this point in the history

  6. Merge pull request #741 from google:loop_clean 

    PiperOrigin-RevId: 648862224
    maxtext authors committed Jul 2, 2024
    Configuration menu
    Copy the full SHA
    c51ee0f View commit details
    Browse the repository at this point in the history

Commits on Jul 3, 2024

  1. Merge branch 'main' into emergency-restore

    @xuefgu
    xuefgu committed Jul 3, 2024
    Configuration menu
    Copy the full SHA
    1429d6c View commit details
    Browse the repository at this point in the history

  2. Add convergence tests on A3 GPU

    @michelle-yooh
    michelle-yooh committed Jul 3, 2024
    Configuration menu
    Copy the full SHA
    24a62b6 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #740 from google:emergency-restore 

    PiperOrigin-RevId: 648923148
    maxtext authors committed Jul 3, 2024
    Configuration menu
    Copy the full SHA
    410dcad View commit details
    Browse the repository at this point in the history

  4. Merge pull request #746 from google:yooh/convergence-test 

    PiperOrigin-RevId: 648942857
    maxtext authors committed Jul 3, 2024
    Configuration menu
    Copy the full SHA
    290dc66 View commit details
    Browse the repository at this point in the history

  5. Update tile size

    @RissyRan
    RissyRan committed Jul 3, 2024
    Configuration menu
    Copy the full SHA
    4e13ea3 View commit details
    Browse the repository at this point in the history

  6. Merge pull request #747 from google:tile_size 

    PiperOrigin-RevId: 649126836
    maxtext authors committed Jul 3, 2024
    Configuration menu
    Copy the full SHA
    3000396 View commit details
    Browse the repository at this point in the history

  7. Handle cases where memstats are not available for the device. 

    Memstats are not guaranteed to be available and can throw an error or return None. This change will handle both `jaxlib.xla_extension.XlaRuntimeError` if the device is not a PjRt addressable device or `KeyError` if the memstats returns None if they are not available.
    @lukebaumann
    lukebaumann committed Jul 3, 2024
    Configuration menu
    Copy the full SHA
    225f197 View commit details
    Browse the repository at this point in the history

  8. Merge pull request #739 from google:memstats 

    PiperOrigin-RevId: 649176682
    maxtext authors committed Jul 3, 2024
    Configuration menu
    Copy the full SHA
    89edace View commit details
    Browse the repository at this point in the history

  9. Fix validation error for other models

    @RissyRan
    RissyRan committed Jul 3, 2024
    Configuration menu
    Copy the full SHA
    eb5cf82 View commit details
    Browse the repository at this point in the history

Commits on Jul 4, 2024

  1. Merge pull request #750 from google:fix_mega_validation 

    PiperOrigin-RevId: 649248081
    maxtext authors committed Jul 4, 2024
    Configuration menu
    Copy the full SHA
    b83a7a4 View commit details
    Browse the repository at this point in the history

Commits on Jul 8, 2024

  1. Fix decode.py to also use first_token from prefill_call

    @vipannalla
    vipannalla committed Jul 8, 2024
    Configuration menu
    Copy the full SHA
    85c105e View commit details
    Browse the repository at this point in the history

  2. Merge pull request #756 from google:decode_prefill_fix 

    PiperOrigin-RevId: 650409674
    maxtext authors committed Jul 8, 2024
    Configuration menu
    Copy the full SHA
    1b4cd15 View commit details
    Browse the repository at this point in the history

Commits on Jul 9, 2024

  1. Add moe perf number

    @RissyRan
    RissyRan committed Jul 9, 2024
    Configuration menu
    Copy the full SHA
    7534e5e View commit details
    Browse the repository at this point in the history

  2. Merge pull request #751 from google:moe_result 

    PiperOrigin-RevId: 650424212
    maxtext authors committed Jul 9, 2024
    Configuration menu
    Copy the full SHA
    5bc4029 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #730 from google:smarter-sharding-parsing 

    PiperOrigin-RevId: 650652469
    maxtext authors committed Jul 9, 2024
    Configuration menu
    Copy the full SHA
    a3d3260 View commit details
    Browse the repository at this point in the history

  4. move num_experts pyconfig assertion to fix tests

    @gobbleturk
    gobbleturk committed Jul 9, 2024
    Configuration menu
    Copy the full SHA
    2cebd1e View commit details
    Browse the repository at this point in the history

Commits on Jul 10, 2024

  1. Merge pull request #760 from google:mattdavidow-fix-assert 

    PiperOrigin-RevId: 650824599
    maxtext authors committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    d0c70e9 View commit details
    Browse the repository at this point in the history

  2. Cast type for inputs before kernel call

    @RissyRan
    RissyRan committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    a26715e View commit details
    Browse the repository at this point in the history

  3. Merge pull request #759 from google:moe_weight_type 

    PiperOrigin-RevId: 650851189
    maxtext authors committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    84d8d05 View commit details
    Browse the repository at this point in the history

  4. Move sharding overrides to models/ directory. 

    PiperOrigin-RevId: 650994392
    @golechwierowicz
    golechwierowicz authored and maxtext authors committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    b506266 View commit details
    Browse the repository at this point in the history

  5. Enable quantization for MoE Matmul implementation

    @RissyRan
    RissyRan committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    6a0e570 View commit details
    Browse the repository at this point in the history

  6. Merge pull request #757 from google:moe_quantization 

    PiperOrigin-RevId: 651070063
    maxtext authors committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    704ab1c View commit details
    Browse the repository at this point in the history

  7. Integrate and test Goodput Monitor with MaxText

    @dipannita08
    dipannita08 committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    f27f70a View commit details
    Browse the repository at this point in the history

  8. Merge pull request #749 from google:integrate-goodput-monitor 

    PiperOrigin-RevId: 651091628
    maxtext authors committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    0af4ee2 View commit details
    Browse the repository at this point in the history

  9. Adding Tokens/s/device to the log.

    @tonyjohnchen
    tonyjohnchen committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    cae41ec View commit details
    Browse the repository at this point in the history

  10. Merge pull request #761 from google:token_per_s 

    PiperOrigin-RevId: 651113335
    maxtext authors committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    5f15eba View commit details
    Browse the repository at this point in the history

  11. Adding support for mixed precision quantization configs.

    @singh-mitali
    singh-mitali committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    8fb12fc View commit details
    Browse the repository at this point in the history

Commits on Jul 11, 2024

  1. Merge pull request #762 from google:msingh-mp 

    PiperOrigin-RevId: 651422435
    maxtext authors committed Jul 11, 2024
    Configuration menu
    Copy the full SHA
    8143e2d View commit details
    Browse the repository at this point in the history