Skip to content

Commit

Permalink
Print some information about device we are running on (#272)
Browse files Browse the repository at this point in the history
* Print some information about device we are running on

On Darwin, use sysctl to get machine name, on Linux, use `/proc/cpuinfo` and if CUDA device is used use `torch.cuda.get_device_name`

Also, change this logger.info to `print` as right now there are no way to control log levels from the generate script
  • Loading branch information
malfet committed Jul 17, 2024
1 parent b0854dd commit d5b21af
Showing 1 changed file with 25 additions and 1 deletion.
26 changes: 25 additions & 1 deletion generate.py
Original file line number Diff line number Diff line change
Expand Up @@ -319,6 +319,30 @@ def encode_tokens(tokenizer, string, bos=True, device="cpu"):
B_INST, E_INST = "[INST]", "[/INST]"


def get_device_info(name: str) -> str:
import platform
from subprocess import check_output

if name == "cpu":
if platform.system() == "Darwin":
return (
check_output(["sysctl", "-n", "machdep.cpu.brand_string"])
.decode("utf-8")
.strip()
)
if platform.system() == "Linux":
return (
check_output(
["sed", "-nr", "s/^model name\\s+: (.*)$/\\1/p", "/proc/cpuinfo"]
)
.decode("utf-8")
.split("\n")[0]
)
if name == "cuda":
return torch.cuda.get_device_name(0)
return ""


def _main(
builder_args: BuilderArgs,
speculative_builder_args: BuilderArgs,
Expand All @@ -341,7 +365,7 @@ def _main(
# # only print on rank 0
# print = lambda *args, **kwargs: None

logging.info(f"Using device={builder_args.device}")
print(f"Using device={builder_args.device} {get_device_info(builder_args.device)}")
set_precision(builder_args.precision)
is_speculative = speculative_builder_args.checkpoint_path is not None

Expand Down

0 comments on commit d5b21af

Please sign in to comment.