Fix tests for llama

bigscience-workshop · Aug 8, 2023 · 070e971 · 070e971
1 parent 17cae64
commit 070e971
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/.github/workflows/run-tests.yaml b/.github/workflows/run-tests.yaml
@@ -40,6 +40,7 @@ jobs:
           export MODEL_NAME=${{ matrix.model }}
           export REF_NAME=${{ matrix.model }}
           export ADAPTER_NAME=${{ matrix.model == 'bigscience/bloom-560m' && 'artek0chumak/bloom-560m-safe-peft' || '' }}
+          export TENSOR_PARALLEL_ARGS=${{ matrix.model == 'bigscience/bloom-560m' && '--tensor_parallel_devices cpu cpu' || '' }}
 
           python -m petals.cli.run_dht --identity tests/test.id --host_maddrs /ip4/127.0.0.1/tcp/31337 &> bootstrap.log &
           BOOTSTRAP_PID=$!
@@ -67,7 +68,7 @@ jobs:
           SERVER3_PID=$!
 
           python -m petals.cli.run_server $MODEL_NAME --num_blocks 5 --torch_dtype float32 \
-            --tensor_parallel_devices cpu cpu \
+            $TENSOR_PARALLEL_ARGS \
             --initial_peers $INITIAL_PEERS --throughput eval &> server4.log &
           SERVER4_PID=$!
 

diff --git a/tests/test_server_stats.py b/tests/test_server_stats.py
@@ -10,7 +10,7 @@
 
 
 @pytest.mark.forked
-def test_server_info(block_from: int = 22, block_to: int = 24, max_length: int = 100, max_length2: int = 50):
+def test_server_info(block_from: int = 6, block_to: int = 8, max_length: int = 100, max_length2: int = 50):
     config = AutoDistributedConfig.from_pretrained(MODEL_NAME)
     dht = hivemind.DHT(initial_peers=INITIAL_PEERS, client_mode=True, start=True)
     blocks1 = RemoteSequential(config, dht=dht, start_block=block_from, end_block=block_to)