merge

Author: priyakasimbeg

algoperf/workloads/librispeech_deepspeech/librispeech_jax/workload.py

@@ -44,6 +44,7 @@ def init_model_fn(
     fake_input_batch = [np.zeros((2, *x), jnp.float32) for x in input_shape]
 
     model_init_fn = jax.jit(functools.partial(self._model.init, train=False))
+    # model_init_fn = functools.partial(self._model.init, train=False)
 
     params_rng, dropout_rng = jax.random.split(rng, 2)
     variables = model_init_fn({'params': params_rng, 'dropout': dropout_rng},

pyproject.toml

@@ -99,7 +99,7 @@ wmt = ["sentencepiece==0.2.0", "tensorflow-text==2.18.0"]
 
 # Frameworks
 jax_core_deps = [
-  "flax==0.8.4",
+  "flax==0.10.4",
   "optax==0.2.2",
   "chex==0.1.86",
   "ml_dtypes==0.4.1",

reference_algorithms/paper_baselines/adamw/jax/submission.py

@@ -75,7 +75,6 @@ def _loss_fn(params):
         spec.ForwardPassMode.TRAIN,
         rng,
         update_batch_norm=True,)
-    jax.debug.print("logits: {logits}", logits=logits)
     loss_dict = workload.loss_fn(
         label_batch=batch['targets'],
         logits_batch=logits,
@@ -222,7 +221,7 @@ def get_batch_size(workload_name):
   elif workload_name == 'librispeech_conformer':
     return 256
   elif workload_name == 'librispeech_deepspeech':
-    return 256
+    return 16
   elif workload_name == 'ogbg':
     return 512
   elif workload_name == 'wmt':

scoring/utils/slurm/make_job_config.py

@@ -0,0 +1,118 @@
+import json
+import os
+
+from absl import app
+from absl import flags
+import jax
+
+SUBMISSION_PATH = 'prize_qualification_baselines/self_tuning/jax_nadamw_full_budget.py'
+EXPERIMENT_DIR = 'submissions/rolling_leaderboard/self_tuning/baseline'
+TUNING_SEARCH_SPACE = None
+FRAMEWORK = 'jax'
+TUNING_RULESET = 'self'
+
+flags.DEFINE_string('submission_path',
+                    SUBMISSION_PATH,
+                    'Path to submission module.')
+flags.DEFINE_string('tuning_search_space',
+                    TUNING_SEARCH_SPACE,
+                    'Path to tuning search space for submission module.')
+flags.DEFINE_string('experiment_dir',
+                    EXPERIMENT_DIR,
+                    'Path to experiment dir where logs will be saved.')
+flags.DEFINE_enum(
+    'framework',
+    FRAMEWORK,
+    enum_values=['jax', 'pytorch'],
+    help='Can be either pytorch or jax.')
+flags.DEFINE_integer('seed', 0, 'RNG seed to to generate study seeds from.')
+flags.DEFINE_enum(
+    'tuning_ruleset',
+    TUNING_RULESET,
+    enum_values=['external', 'self'],
+    help='Which tuning ruleset to score this submission on. Can be external or self.'
+)
+
+FLAGS = flags.FLAGS
+
+MIN_INT = -2**(31)
+MAX_INT = 2**(31) - 1
+NUM_TUNING_TRIALS = 5  # For external tuning ruleset
+NUM_STUDIES = 3
+
+WORKLOADS = {
+    "imagenet_resnet": {"dataset": "imagenet"},
+    "imagenet_vit": {"dataset": "imagenet"},
+    "fastmri": {"dataset": "fastmri"},
+    "ogbg": {"dataset": "ogbg"},
+    "wmt": {"dataset": "wmt"},
+    "librispeech_deepspeech": {"dataset": "librispeech"},
+    "criteo1tb": {"dataset": "criteo1tb"},
+    "librispeech_conformer": {"dataset": "librispeech"}
+}
+
+
+def main(_):
+  workloads = WORKLOADS.keys()
+  key = jax.random.key(FLAGS.seed)
+
+  jobs = []
+
+  for workload in workloads:
+    # Fold in hash(workload) mod(max(uint32))
+    workload_key = jax.random.fold_in(key, hash(workload) % (2**32 - 1))
+    for study_index in range(NUM_STUDIES):
+      study_key = jax.random.fold_in(workload_key, study_index)
+      if FLAGS.tuning_ruleset == 'external':
+        for hparam_index in range(NUM_TUNING_TRIALS):
+          run_key = jax.random.fold_in(study_key, hparam_index)
+          seed = jax.random.randint(run_key, (1,), MIN_INT, MAX_INT)[0].item()
+          print(seed)
+          # Add job
+          job = {}
+          study_dir = os.path.join(FLAGS.experiment_dir, f"study_{study_index}")
+          job['framework'] = FLAGS.framework
+          job['workload'] = workload
+          job['dataset'] = WORKLOADS[workload]['dataset']
+          job['submission_path'] = FLAGS.submission_path
+          job['experiment_dir'] = study_dir
+          job['rng_seed'] = seed
+          job['tuning_ruleset'] = FLAGS.tuning_ruleset
+          job['num_tuning_trials'] = NUM_TUNING_TRIALS
+          job['hparam_start_index'] = hparam_index
+          job['hparam_end_index'] = hparam_index + 1
+          job['tuning_search_space'] = FLAGS.tuning_search_space
+          job['tuning_ruleset'] = FLAGS.tuning_ruleset
+          jobs.append(job)
+          print(job)
+
+      else:
+        run_key = study_key
+        seed = jax.random.randint(run_key, (1,), MIN_INT, MAX_INT)[0].item()
+        print(seed)
+        # Add job
+        job = {}
+        study_dir = os.path.join(FLAGS.experiment_dir, f"study_{study_index}")
+        job['framework'] = FLAGS.framework
+        job['workload'] = workload
+        job['dataset'] = WORKLOADS[workload]['dataset']
+        job['submission_path'] = FLAGS.submission_path
+        job['experiment_dir'] = study_dir
+        job['rng_seed'] = seed
+        job['tuning_ruleset'] = FLAGS.tuning_ruleset
+        job['num_tuning_trials'] = 1
+
+        jobs.append(job)
+        print(job)
+
+  # Convert job array to dict with job indices
+  job_dict = {}
+  for i, job in enumerate(jobs):
+    job_dict[f"{i}"] = job
+
+  with open('config.json', 'w') as f:
+    json.dump(job_dict, f, indent=4)
+
+
+if __name__ == '__main__':
+  app.run(main)

submission_runner.py

@@ -636,20 +636,22 @@ def score_submission_on_workload(workload: spec.Workload,
         tuning_search_space[hi] = hyperparameters
 
       with profiler.profile('Train'):
-        timing, metrics = train_once(workload, workload_name,
-                                     global_batch_size,
-                                     global_eval_batch_size,
-                                     data_dir, imagenet_v2_data_dir,
-                                     init_optimizer_state,
-                                     update_params, data_selection,
-                                     prepare_for_eval,
-                                     hyperparameters,
-                                     rng_seed,
-                                     rng,
-                                     profiler,
-                                     max_global_steps,
-                                     tuning_dir_name,
-                                     save_checkpoints=save_checkpoints,)
+        with jax.profiler.trace("/logs/tensorboard"):
+          print('profiling!')
+          timing, metrics = train_once(workload, workload_name,
+                                      global_batch_size,
+                                      global_eval_batch_size,
+                                      data_dir, imagenet_v2_data_dir,
+                                      init_optimizer_state,
+                                      update_params, data_selection,
+                                      prepare_for_eval,
+                                      hyperparameters,
+                                      rng_seed,
+                                      rng,
+                                      profiler,
+                                      max_global_steps,
+                                      tuning_dir_name,
+                                      save_checkpoints=save_checkpoints,)
       all_timings[hi] = timing
       all_metrics[hi] = metrics
       logging.info(f'Tuning trial {hi + 1}/{num_tuning_trials}')