Merge pull request #265 from mila-iqia/staging

Staging

.gitignore

@@ -34,6 +34,7 @@ scripts/article/xpu/
 dependencies/
 benchmarks/recursiongfn/gflownet
 benchmarks/recursiongfn/logs/
+benchmarks/llm/tune/
 
 scripts/inventory.yaml
 output/

benchmarks/cleanrl_jax/Makefile

@@ -0,0 +1,31 @@
+# Use global base if possible
+ifndef MILABENCH_BASE
+	MILABENCH_BASE="base"
+endif
+
+export MILABENCH_BASE
+
+BENCH_NAME=cleanrl_jax
+MILABENCH_CONFIG=dev.yaml
+MILABENCH_ARGS=--config $(MILABENCH_CONFIG) --base $(MILABENCH_BASE)
+
+all:
+	install prepare single gpus nodes
+
+install:
+	milabench install $(MILABENCH_ARGS) --force
+
+prepare:
+	milabench prepare $(MILABENCH_ARGS)
+
+tests: install prepare
+	milabench run $(MILABENCH_ARGS)
+
+single:
+	milabench run $(MILABENCH_ARGS) --select $(BENCH_NAME)
+
+gpus:
+	milabench run $(MILABENCH_ARGS) --select $(BENCH_NAME)-gpus
+
+nodes:
+	milabench run $(MILABENCH_ARGS) --select $(BENCH_NAME)-nodes

benchmarks/cleanrl_jax/README.md

@@ -0,0 +1,4 @@
+
+# Cleanrl_jax
+
+Rewrite this README to explain what the benchmark is!

benchmarks/cleanrl_jax/benchfile.py

@@ -0,0 +1,31 @@
+from milabench.pack import Package
+
+
+class Cleanrl_jax(Package):
+    # Requirements file installed by install(). It can be empty or absent.
+    base_requirements = "requirements.in"
+
+    # The preparation script called by prepare(). It must be executable,
+    # but it can be any type of script. It can be empty or absent.
+    prepare_script = "prepare.py"
+
+    # The main script called by run(). It must be a Python file. It has to
+    # be present.
+    main_script = "main.py"
+
+    # You can remove the functions below if you don't need to modify them.
+
+    def make_env(self):
+        # Return a dict of environment variables for prepare_script and
+        # main_script.
+        return super().make_env()
+
+    async def install(self):
+        await super().install()  # super() call installs the requirements
+
+    async def prepare(self):
+        await super().prepare()  # super() call executes prepare_script
+
+
+
+__pack__ = Cleanrl_jax

benchmarks/cleanrl_jax/dev.yaml

@@ -0,0 +1,8 @@
+
+cleanrl_jax:
+  inherits: _defaults
+  definition: .
+  install-variant: unpinned
+  install_group: torch
+  plan:
+    method: per_gpu