Update workflow

huggingface · Jan 4, 2024 · c88fe86 · c88fe86
1 parent ec39922
commit c88fe86
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 5 deletions.
diff --git a/.github/workflows/test_trainium_common.yml b/.github/workflows/test_trainium_common.yml
@@ -32,6 +32,8 @@ jobs:
         run: echo "/home/ubuntu/.local/bin" >> $GITHUB_PATH
       - name: Set pip repository pointing to the Neuron repository
         run: pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
+      - name: Update pip
+        run: pip install -U pip
       - name: Install Python dependencies
         run: pip install .[tests,neuronx]
       - name: Run tests on Neuron cores

diff --git a/tests/distributed/test_model_parallelization.py b/tests/distributed/test_model_parallelization.py
@@ -359,11 +359,6 @@ def _parallel_model_matches_original_model(
         static_seed_patcher = create_static_seed_patcher(model.__class__, 42)
         with static_seed_patcher:
             model = accelerator.prepare(model)
-        if xm.get_ordinal() == 0:
-            pass
-            # print(model.gpt_neox.embed_in.weight, orig_model.gpt_neox.embed_in.weight)
-            # print(model.embed_out.weight, orig_model.embed_out.weight)
-            # print(model.gpt_neox.embed_in.weight, model.embed_out.weight)
 
         with torch.no_grad():
             if pp_size == 1: