Code

Draft: Refactor JobSet for Pathways#918

Code

Draft

jiya-zhang wants to merge 65 commits intoapple:mainapple/axlearn:mainfrom jesus-orozco:pathways_trilliumjesus-orozco/axlearn:pathways_trilliumCopy head branch name to clipboard

+255-63

Commits

Commits on Sep 9, 2024

• Adding support for Pathways proxy

  jesus-orozco

  committed
  582329a

  Copy full SHA for 582329a

Commits on Oct 1, 2024

• Update pathways-utils dependency and fix formatting

  jesus-orozco

  committed
  8d3c643

  Copy full SHA for 8d3c643

Commits on Oct 7, 2024

• Move pathways package to its own dependency tree and pin it to a specific tagged version

  jesus-orozco

  authored
  0e61b76

  Copy full SHA for 0e61b76

Commits on Oct 9, 2024

• Relocate pathwaysutils import

  jesus-orozco

  committed
  4260c38

  Copy full SHA for 4260c38

• Create custom jobset for pathways

  jesus-orozco

  committed
  ae80a36

  Copy full SHA for ae80a36

Commits on Oct 16, 2024

• Updates to pathways jobset creation

  jesus-orozco

  committed
  9960027

  Copy full SHA for 9960027

Commits on Oct 17, 2024

• Merge branch 'main' into feature/jax_pathways

  jesus-orozco

  authored
  f02d345

  Copy full SHA for f02d345

Commits on Oct 25, 2024

• Merge branch 'apple:main' into feature/jax_pathways

  jesus-orozco

  authored
  6c3c083

  Copy full SHA for 6c3c083

• Update pathwaysutils source to pypi

  jesus-orozco

  authored
  ac6bcd2

  Copy full SHA for ac6bcd2

Commits on Oct 28, 2024

• Merge branch 'main' into feature/pathways_workload

  jesus-orozco

  authored
  43ac5e4

  Copy full SHA for 43ac5e4

• Merge branch 'apple:main' into feature/jax_pathways

  jesus-orozco

  authored
  05f311d

  Copy full SHA for 05f311d

Commits on Nov 6, 2024

• trillium testing baseline

  jesus-orozco

  committed
  8972933

  Copy full SHA for 8972933

• revert dockerfile for upstream merge

  jesus-orozco

  committed
  a952c6f

  Copy full SHA for a952c6f

• Merge branch 'apple:main' into trillium_testing

  jesus-orozco

  authored
  34571aa

  Copy full SHA for 34571aa

• fixed pdbs 3 for fuji 70b

  jesus-orozco

  committed
  0e8ae86

  Copy full SHA for 0e8ae86

• testing pdbs 3 with 2 v6e-256 slices

  jesus-orozco

  committed
  d9d458a

  Copy full SHA for d9d458a

Commits on Nov 7, 2024

• use maxtext xla flags only

  jesus-orozco

  committed
  73da333

  Copy full SHA for 73da333

• new baseline for pdbs=3 without ffn_dim

  jesus-orozco

  committed
  01bcf8f

  Copy full SHA for 01bcf8f

• try xla sc offload flags

  jesus-orozco

  committed
  b1ba4fe

  Copy full SHA for b1ba4fe

• revert AR + SC offload flags

  jesus-orozco

  committed
  6b24cd2

  Copy full SHA for 6b24cd2

• output jobset to yaml file

  jesus-orozco

  committed
  aa06778

  Copy full SHA for aa06778

• calculate batch size based on flags

  jesus-orozco

  committed
  efabbfd

  Copy full SHA for efabbfd

• enable ffn 3.5 and test pdbs 3 with 4 slices

  jesus-orozco

  committed
  cd556c7

  Copy full SHA for cd556c7

• retry 4 slices with pdbs 3

  jesus-orozco

  committed
  90b0d26

  Copy full SHA for 90b0d26

Commits on Nov 8, 2024

• test xla_tpu_enable_sparse_core_collective_offload_all_reduce

  jesus-orozco

  committed
  4974458

  Copy full SHA for 4974458

• remove xla_tpu_enable_sparse_core_collective_offload_all_reduce

  jesus-orozco

  committed
  caff395

  Copy full SHA for caff395

• dynamic global batch size based on pdbs and slices

  jesus-orozco

  committed
  8277980

  Copy full SHA for 8277980

Commits on Nov 11, 2024

• enable xla_enable_async_all_reduce

  jesus-orozco

  committed
  e23a4d0

  Copy full SHA for e23a4d0

• Merge branch 'main' into feature/jax_pathways

  jesus-orozco

  authored
  1797c52

  Copy full SHA for 1797c52

• Refactor pathways config flag

  jesus-orozco

  committed
  4868815

  Copy full SHA for 4868815

• Merge branch 'apple:main' into feature/pathways_workload

  jesus-orozco

  authored
  fe62afd

  Copy full SHA for fe62afd

• install libtpu nightly

  jesus-orozco

  committed
  27ceea0

  Copy full SHA for 27ceea0

• Merge branch 'apple:main' into trillium_testing

  jesus-orozco

  authored
  5c80fc8

  Copy full SHA for 5c80fc8

Commits on Nov 15, 2024

• custom remat policy for fuji-70b

  jesus-orozco

  committed
  200ac48

  Copy full SHA for 200ac48

• sparscore xla flags and nothing_saveable remat policy

  jesus-orozco

  committed
  3bfda15

  Copy full SHA for 3bfda15

Commits on Nov 18, 2024

• calculate batch size with jax devices

  jesus-orozco

  committed
  51e2e90

  Copy full SHA for 51e2e90

• update remat policy offloading

  jesus-orozco

  committed
  b894304

  Copy full SHA for b894304

• Merge branch 'main' into feature/pathways_workload

  jesus-orozco

  authored
  ce5f1a2

  Copy full SHA for ce5f1a2

• Update axlearn/cloud/gcp/job.py

  Show description for b51e67e
  

  jesus-orozco

  and

  ruomingp

  authored
  b51e67e

  Copy full SHA for b51e67e

• Update job.py with dynamic module imports

  jesus-orozco

  authored
  6537274

  Copy full SHA for 6537274

• Update job.py - remove pathways from dynamic import error message

  jesus-orozco

  authored
  7dbb1b9

  Copy full SHA for 7dbb1b9

Commits on Nov 22, 2024

• Merge remote-tracking branch 'origin/feature/jax_pathways' into pathways_trillium

  jesus-orozco

  committed
  af7c746

  Copy full SHA for af7c746

• pathways jobset updates

  jesus-orozco

  committed
  5704ce5

  Copy full SHA for 5704ce5

• merge trillium changes

  jesus-orozco

  committed
  bf82554

  Copy full SHA for bf82554

Commits on Nov 25, 2024

• Merge branch 'apple:main' into pathways_trillium

  jesus-orozco

  authored
  4f58291

  Copy full SHA for 4f58291

• Install pathwaysutils

  jesus-orozco

  committed
  d206110

  Copy full SHA for d206110

• Disable force eval

  jesus-orozco

  committed
  4152db3

  Copy full SHA for 4152db3

• Launch pathways on trainer_main

  jesus-orozco

  committed
  0c33403

  Copy full SHA for 0c33403

Commits on Dec 5, 2024

• add v6e mesh rules

  jesus-orozco

  committed
  315fa6b

  Copy full SHA for 315fa6b

• pin libtpu version

  jesus-orozco

  committed
  765f64b

  Copy full SHA for 765f64b

• update pathways jobset definition

  jesus-orozco

  committed
  f387807

  Copy full SHA for f387807

Commits on Dec 6, 2024

• Merge branch 'apple:main' into pathways_trillium

  jesus-orozco

  authored
  735ff37

  Copy full SHA for 735ff37

Commits on Dec 9, 2024

• dump xla flags to gcs

  jesus-orozco

  committed
  2e0dd4b

  Copy full SHA for 2e0dd4b

Commits on Dec 12, 2024

• Refactor jobset to align with new pathways structure

  jesus-orozco

  committed
  e9d8341

  Copy full SHA for e9d8341

• Apply formatting

  jesus-orozco

  committed
  357a322

  Copy full SHA for 357a322

Commits on Jan 10, 2025

• refactor pathways jobset to new spec

  jesus-orozco

  committed
  89ac1ea

  Copy full SHA for 89ac1ea

Commits on Jan 14, 2025

• Rebase to axlearn main

  jiya-zhang

  committed
  d63dd6a

  Copy full SHA for d63dd6a

• revert changes to gitignore and remove dockerignore

  jesus-orozco

  committed
  4393453

  Copy full SHA for 4393453

• revert changes to dependencies

  jesus-orozco

  committed
  9ba599b

  Copy full SHA for 9ba599b

• revert changes to trainer and model configs

  jesus-orozco

  committed
  ca5c883

  Copy full SHA for ca5c883

• remove unnecessary updates to gke tpu job for pathways workloads

  jesus-orozco

  committed
  c7bc5df

  Copy full SHA for c7bc5df

• Revert updates to fuji model config

  jesus-orozco

  committed
  a0bf9df

  Copy full SHA for a0bf9df

Commits on Jan 15, 2025

• Update pathways container specs for gke tpu job

  jesus-orozco

  committed
  9f19159

  Copy full SHA for 9f19159

Commits on Jan 16, 2025

• Update gke tpu job to bypass jobset coordinator

  jesus-orozco

  committed
  8f71508

  Copy full SHA for 8f71508

Commits on Jan 31, 2025

• Swap pathways check for tpu gke job config flag

  jesus-orozco

  committed
  f467064

  Copy full SHA for f467064