-

benchmarks/lightning/main.py

@@ -55,6 +55,8 @@ def prepare_voir():
     return observer, bench_monitor
 
 def main():
+    from habana_frameworks.torch import hpu; hpu.init()
+
     rank = int(os.getenv("RANK", 0))
     world_size = int(os.getenv("WORLD_SIZE", 1))
     local_world_size = int(os.getenv("LOCAL_WORLD_SIZE", 1))

docker/Makefile

@@ -3,7 +3,7 @@
 
 # bench = rlhf-single
 bench = "lightning"
-args = "--exclude lightning-gpus"
+lazy = 0
 
 hpu:
 	git add --all
@@ -13,5 +13,5 @@ hpu:
 	# docker system prune -a -f
 	# docker image prune -a -f
 	docker build --build-arg BENCH=$(bench) --build-arg CACHEBUST=`git rev-parse hpu` -f Dockerfile-hpu -t dockerfile-hpu . 
-	docker run   -it   --runtime=habana  -e PT_HPU_LAZY_MODE=0 -e HABANA_VISIBLE_DEVICES=all   -e OMPI_MCA_btl_vader_single_copy_mechanism=none --shm-size 50G   --cap-add=sys_nice   --net=host   dockerfile-hpu:latest bash -c '. $$MILABENCH_VENV/bin/activate && milabench install --use-current-env --select $(bench) && pip uninstall torch torchvision torchaudio -y  &&  milabench prepare --use-current-env --select $(bench) && milabench run --use-current-env $(args) --select $(bench)'
+	docker run   -it   --runtime=habana  -e PT_HPU_LAZY_MODE=$(lazy) -e HABANA_VISIBLE_DEVICES=all   -e OMPI_MCA_btl_vader_single_copy_mechanism=none --shm-size 50G   --cap-add=sys_nice   --net=host   dockerfile-hpu:latest bash -c '. $$MILABENCH_VENV/bin/activate && milabench install --use-current-env --select $(bench) && pip uninstall torch torchvision torchaudio -y  &&  milabench prepare --use-current-env --select $(bench) && milabench run --use-current-env $(args) --exclude lightning-gpus --select $(bench)'