Stas Bekman stas00

Connect via SSH to a Slurm compute job that runs as Enroot container

Being able to SSH directly into a compute job has the advantage of using all remote development tools such as using your IDE's debugger also for GPU jobs (VSCode, PyCharm, ...).

Slurm: Scheduling system that many HPC clusters use
Enroot: Container system like Docker for NVIDIA GPUs

General problem:

Locate the section for your github remote in the .git/config file. It looks like this:

[remote "origin"]
	fetch = +refs/heads/*:refs/remotes/origin/*
	url = [email protected]:joyent/node.git

Now add the line fetch = +refs/pull/*/head:refs/remotes/origin/pr/* to this section. Obviously, change the github url to match your project's URL. It ends up looking like this:

	# This is a modified version of TRL's `SFTTrainer` example (https://github.com/huggingface/trl/blob/main/examples/scripts/sft_trainer.py),
	# adapted to run with DeepSpeed ZeRO-3 and Mistral-7B-V1.0. The settings below were run on 1 node of 8 x A100 (80GB) GPUs.
	#
	# Usage:
	# - Install the latest transformers & accelerate versions: `pip install -U transformers accelerate`
	# - Install deepspeed: `pip install deepspeed==0.9.5`
	# - Install TRL from main: pip install git+https://github.com/huggingface/trl.git
	# - Clone the repo: git clone github.com/huggingface/trl.git
	# - Copy this Gist into trl/examples/scripts
	# - Run from root of trl repo with: accelerate launch --config_file=examples/accelerate_configs/deepspeed_zero3.yaml --gradient_accumulation_steps 8 examples/scripts/sft_trainer.py

	import torch
	from torch.utils.flop_counter import FlopCounterMode
	from triton.testing import do_bench

	def get_flops_achieved(f):
	flop_counter = FlopCounterMode(display=False)
	with flop_counter:
	f()
	total_flops = flop_counter.get_total_flops()
	ms_per_iter = do_bench(f)

	# Reference
	# =========
	# Why does deep and cheap learning work so well?∗
	# Henry W. Lin, Max Tegmark, and David Rolnick
	# Dept. of Physics, Harvard University, Cambridge, MA 02138
	# Dept. of Physics, Massachusetts Institute of Technology, Cambridge, MA 02139 and Dept. of Mathematics, Massachusetts Institute of Technology, Cambridge, MA 02139
	# Here the x input takes two numbers and produces the multiplication as output
	import numpy as np

	# multiplication approximator

	'''
	Memory profiling utilities
	'''
	import gc
	import inspect
	import linecache
	import os.path
	import sys
	import time
	import threading

	#!/usr/bin/env bash

	git remote -v \| grep fetch \| grep github \| \
	while read remote url _; do
	if ! git config --get-all "remote.$remote.fetch" \| grep -q refs/pull
	then
	git config --add "remote.$remote.fetch" \
	'+refs/pull//head:refs/remotes/'"$remote"'/pull/'
	fi
	done