Merge branch 'lm_workload' of github.com:mlcommons/algorithmic-efficiency into lm_workload

priyakasimbeg · priyakasimbeg · commit 42d91aca751f · 2025-11-08T01:04:59.000Z
diff --git a/algoperf/workloads/finewebedu_lm/__init__.py b/algoperf/workloads/finewebedu_lm/__init__.py
diff --git a/algoperf/workloads/finewebedu_lm/finewebedu_lm_jax/__init__.py b/algoperf/workloads/finewebedu_lm/finewebedu_lm_jax/__init__.py
diff --git a/algoperf/workloads/finewebedu_lm/finewebedu_lm_jax/models.py b/algoperf/workloads/finewebedu_lm/finewebedu_lm_jax/models.py
diff --git a/algoperf/workloads/finewebedu_lm/finewebedu_lm_jax/workload.py b/algoperf/workloads/finewebedu_lm/finewebedu_lm_jax/workload.py
@@ -6,12 +6,12 @@
 import jax.numpy as jnp
 
 from algoperf import jax_sharding_utils, param_utils, spec
-from algoperf.workloads.lm.input_pipeline import get_data_iter
-from algoperf.workloads.lm.lm_jax.nanodo_model import (
+from algoperf.workloads.finewebedu_lm.finewebedu_lm_jax.models import (
   ModelConfig,
   TransformerDo,
 )
-from algoperf.workloads.lm.workload import BaseLmWorkload
+from algoperf.workloads.finewebedu_lm.input_pipeline import get_data_iter
+from algoperf.workloads.finewebedu_lm.workload import BaseLmWorkload
 
 
 class LmWorkload(BaseLmWorkload):
diff --git a/algoperf/workloads/finewebedu_lm/finewebedu_lm_pytorch/__init__.py b/algoperf/workloads/finewebedu_lm/finewebedu_lm_pytorch/__init__.py
diff --git a/algoperf/workloads/finewebedu_lm/finewebedu_lm_pytorch/models.py b/algoperf/workloads/finewebedu_lm/finewebedu_lm_pytorch/models.py
diff --git a/algoperf/workloads/finewebedu_lm/finewebedu_lm_pytorch/workload.py b/algoperf/workloads/finewebedu_lm/finewebedu_lm_pytorch/workload.py
@@ -10,12 +10,12 @@
 from torch.nn.parallel import DistributedDataParallel as DDP
 
 from algoperf import param_utils, pytorch_utils, spec
-from algoperf.workloads.lm.input_pipeline import get_data_iter
-from algoperf.workloads.lm.lm_pytorch.plainlm_model import (
+from algoperf.workloads.finewebedu_lm.finewebedu_lm_pytorch.models import (
   ModelConfig,
   Transformer,
 )
-from algoperf.workloads.lm.workload import BaseLmWorkload
+from algoperf.workloads.finewebedu_lm.input_pipeline import get_data_iter
+from algoperf.workloads.finewebedu_lm.workload import BaseLmWorkload
 
 USE_PYTORCH_DDP, RANK, DEVICE, N_GPUS = pytorch_utils.pytorch_setup()
 
diff --git a/algoperf/workloads/finewebedu_lm/input_pipeline.py b/algoperf/workloads/finewebedu_lm/input_pipeline.py
diff --git a/algoperf/workloads/finewebedu_lm/workload.py b/algoperf/workloads/finewebedu_lm/workload.py
diff --git a/algoperf/workloads/ogbg/workload.py b/algoperf/workloads/ogbg/workload.py
@@ -92,7 +92,7 @@ def max_allowed_runtime_sec(self) -> int:
 
   @property
   def eval_period_time_sec(self) -> int:
-    return 452 # approx 25 evals
+    return 452  # approx 25 evals
 
   def _build_input_queue(
     self,
diff --git a/algoperf/workloads/workloads.py b/algoperf/workloads/workloads.py
@@ -113,7 +113,10 @@
     'workload_path': 'librispeech_deepspeech/librispeech',
     'workload_class_name': 'LibriSpeechDeepSpeechNormAndSpecAugWorkload',
   },
-  'lm': {'workload_path': 'lm/lm', 'workload_class_name': 'LmWorkload'},
+  'finewebedu_lm': {
+    'workload_path': 'finewebedu_lm/finewebedu_lm',
+    'workload_class_name': 'LmWorkload',
+  },
   'mnist': {
     'workload_path': 'mnist/mnist',
     'workload_class_name': 'MnistWorkload',
@@ -153,7 +156,7 @@
   'imagenet_vit',
   'librispeech_conformer',
   'librispeech_deepspeech',
-  'lm',
+  'finewebedu_lm',
   'ogbg',
   'wmt',
 ]
diff --git a/algorithms/archived_paper_baselines/adamw/pytorch/submission.py b/algorithms/archived_paper_baselines/adamw/pytorch/submission.py
@@ -189,8 +189,8 @@ def get_batch_size(workload_name):
     return 128
   elif workload_name == 'mnist':
     return 16
-  elif workload_name == 'lm':
-    return 4
+  elif workload_name == 'finewebedu_lm':
+    return 64
   else:
     raise ValueError(f'Unsupported workload name: {workload_name}.')
 
diff --git a/algorithms/archived_paper_baselines/nesterov/jax/submission.py b/algorithms/archived_paper_baselines/nesterov/jax/submission.py
@@ -292,8 +292,8 @@ def get_batch_size(workload_name):
     return 16
   elif workload_name == 'cifar':
     return 128
-  elif workload_name == 'lm':
-    return 8
+  elif workload_name == 'finewebedu_lm':
+    return 64
   else:
     raise ValueError(f'Unsupported workload name: {workload_name}.')
 
diff --git a/algorithms/baselines/external_tuning/jax_nadamw_full_budget.py b/algorithms/baselines/external_tuning/jax_nadamw_full_budget.py
@@ -394,7 +394,7 @@ def get_batch_size(workload_name):
     return 512
   elif workload_name == 'wmt':
     return 128
-  elif workload_name == 'lm':
+  elif workload_name == 'finewebedu_lm':
     return 64
   elif workload_name == 'mnist':
     return 16
diff --git a/algorithms/baselines/external_tuning/pytorch_nadamw_full_budget.py b/algorithms/baselines/external_tuning/pytorch_nadamw_full_budget.py
@@ -372,7 +372,7 @@ def get_batch_size(workload_name):
     return 128
   elif workload_name == 'mnist':
     return 16
-  elif workload_name == 'lm':
+  elif workload_name == 'finewebedu_lm':
     return 64
   else:
     raise ValueError(f'Unsupported workload name: {workload_name}.')
diff --git a/algorithms/target_setting_algorithms/fineweb_edu_lm/jax_nadamw_target_setting.py b/algorithms/target_setting_algorithms/fineweb_edu_lm/jax_nadamw_target_setting.py
@@ -395,7 +395,7 @@ def get_batch_size(workload_name):
     return 512
   elif workload_name == 'wmt':
     return 128
-  elif workload_name == 'lm':
+  elif workload_name == 'finewebedu_lm':
     return 64
   elif workload_name == 'mnist':
     return 16
diff --git a/algorithms/target_setting_algorithms/fineweb_edu_lm/pytorch_nadamw_target_setting.py b/algorithms/target_setting_algorithms/fineweb_edu_lm/pytorch_nadamw_target_setting.py
@@ -373,7 +373,7 @@ def get_batch_size(workload_name):
     return 128
   elif workload_name == 'mnist':
     return 16
-  elif workload_name == 'lm':
+  elif workload_name == 'finewebedu_lm':
     return 64
   else:
     raise ValueError(f'Unsupported workload name: {workload_name}.')
diff --git a/scoring/performance_profile.py b/scoring/performance_profile.py
@@ -71,7 +71,7 @@
   'wer',
   'l1_loss',
   'loss',
-  'ppl'
+  'ppl',
 ]
 
 MAX_EVAL_METRICS = ['mean_average_precision', 'ssim', 'accuracy', 'bleu']
diff --git a/scoring/utils/workload_metadata_external_tuning.json b/scoring/utils/workload_metadata_external_tuning.json
@@ -31,7 +31,7 @@
     "max_steps": 80000,
     "dataset": "librispeech"
   },
-  "lm" : {
+  "finewebedu_lm" : {
     "max_steps": 55000,
     "dataset":"fineweb_edu_10B"
   }
diff --git a/submission_runner.py b/submission_runner.py
@@ -266,7 +266,7 @@ def train_once(
         'librispeech_deepspeech',
         'ogbg',
         'wmt',
-        'lm'
+        'finewebedu_lm',
       ]
       base_workload = workloads.get_base_workload_name(workload_name)
       if base_workload in compile_error_workloads:
@@ -784,7 +784,7 @@ def main(_):
     os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:256'
 
   if FLAGS.framework == 'pytorch':
-    limit_tf_threads = base_workload != 'lm'
+    limit_tf_threads = base_workload != 'finewebedu_lm'
     pytorch_init(
       USE_PYTORCH_DDP, RANK, profiler, limit_tf_threads=limit_tf_threads
     )
@@ -803,7 +803,7 @@ def main(_):
     'librispeech_deepspeech',
     'imagenet_vit',
     'criteo1tb',
-    'lm',
+    'finewebedu_lm',
   ]:
     os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.80'
 
diff --git a/tests/modeldiffs/lm/compare.py b/tests/modeldiffs/lm/compare.py
@@ -26,28 +26,28 @@
 from absl.testing import absltest, parameterized
 
 # Import JAX implementation
-from algoperf.workloads.lm.lm_jax.nanodo_model import (
+from algoperf.workloads.finewebedu_lm.finewebedu_lm_jax.models import (
   CausalAttn,
   Mlp,
   TBlock,
   TransformerDo,
   apply_rope,
   init_rope,
 )
-from algoperf.workloads.lm.lm_jax.nanodo_model import (
+from algoperf.workloads.finewebedu_lm.finewebedu_lm_jax.models import (
   ModelConfig as JaxModelConfig,
 )
 
 # Import PyTorch implementation
-from algoperf.workloads.lm.lm_pytorch.plainlm_model import (
+from algoperf.workloads.finewebedu_lm.finewebedu_lm_pytorch.models import (
   MLP,
   Attention,
   Block,
   Transformer,
   apply_rotary_emb_complex_like,
   precompute_freqs_cis,
 )
-from algoperf.workloads.lm.lm_pytorch.plainlm_model import (
+from algoperf.workloads.finewebedu_lm.finewebedu_lm_pytorch.models import (
   ModelConfig as PyTorchModelConfig,
 )
 
