Skip to content

Commit

Permalink
Fix llm requirements
Browse files Browse the repository at this point in the history
* rename huggingface token to MILABENCH_* to automatically forward the env var to a remote in such cases
  • Loading branch information
satyaog committed Sep 5, 2024
1 parent 0b1320a commit 39222b8
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 7 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/tests_unit.yml
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ jobs:
- name: tests
env:
HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN}}
MILABENCH_HF_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN}}
run: |
source $(poetry env info -p)/bin/activate
coverage run --source=milabench -m pytest --ignore=tests/integration tests/
Expand Down
9 changes: 8 additions & 1 deletion benchmarks/llm/prepare.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,18 +55,25 @@ def generate_model(
model_parallel_size=1
):
try:
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "12355"

torch.distributed.init_process_group(rank=rank, world_size=model_parallel_size)
fairscale.nn.model_parallel.initialize.initialize_model_parallel(model_parallel_size)

conn.send(os.getpid())
while not conn.poll():
time.sleep(0.1)
conn.recv()

params = json.loads(params_path.read_text())
model = llama.model.Transformer(ModelArgs(**params))
torch.save(model.state_dict(), params_path.with_name(f"consolidated.{rank:02}.pth"))

except Exception as e:
conn.send(e)
raise

finally:
conn.close()

Expand Down Expand Up @@ -101,7 +108,7 @@ def main():
config = OmegaConf.merge(base, cli)

repo_id = config["repo_id"]
hf_token = os.getenv("HUGGING_FACE_TOKEN", None)
hf_token = os.getenv("MILABENCH_HF_TOKEN", None)
output_dir = config["checkpointer"]["output_dir"]

ignore_patterns = ["*.safetensors", "*consolidated.*.pth"]
Expand Down
40 changes: 35 additions & 5 deletions benchmarks/llm/requirements.cuda.txt

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 39222b8

Please sign in to comment.