Update app.py

app.py

@@ -220,6 +220,13 @@ arg1 = '7b'
 arg2 = 32
 arg3 = 0.0002
 
+device_map = "auto"
+world_size = int(os.environ.get("WORLD_SIZE", 1))
+ddp = world_size != 1
+if ddp:
+    device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)}
+    GRADIENT_ACCUMULATION_STEPS = GRADIENT_ACCUMULATION_STEPS // world_size
+
 MICRO_BATCH_SIZE = int(arg2)
 BATCH_SIZE = 64
 size = arg1