gpt-neox

1
# Copyright (c) 2024, EleutherAI
2
# This file is based on code by the authors denoted below and has been modified from its original version.
3
#
4
# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
5
#
6
# Licensed under the Apache License, Version 2.0 (the "License");
7
# you may not use this file except in compliance with the License.
8
# You may obtain a copy of the License at
9
#
10
#     http://www.apache.org/licenses/LICENSE-2.0
11
#
12
# Unless required by applicable law or agreed to in writing, software
13
# distributed under the License is distributed on an "AS IS" BASIS,
14
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
# See the License for the specific language governing permissions and
16
# limitations under the License.
17

18
"""Evaluation tasks - modified from https://github.com/EleutherAI/lm-evaluation-harness"""
19
import os
20
import sys
21

22
sys.path.append(
23
    os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir))
24
)
25
from megatron.training import forward_step
26
from megatron.utils import setup_for_inference_or_eval, init_wandb
27
from megatron.logging import tb_wandb_log
28
from eval_tasks import run_eval_harness
29
from pprint import pprint
30
from datetime import datetime
31
import json
32

33

34
def main(input_args=None, overwrite_values=None):
35
    model, neox_args = setup_for_inference_or_eval(
36
        use_cache=False, input_args=input_args, overwrite_values=overwrite_values
37
    )
38
    results = run_eval_harness(
39
        model,
40
        forward_step,
41
        neox_args,
42
        eval_tasks=neox_args.eval_tasks,
43
        bootstrap_iters=10000,
44
    )
45
    if neox_args.rank == 0:
46
        init_wandb(neox_args=neox_args)
47
        # log to wandb
48
        for k, v in results["results"].items():
49
            if isinstance(v, dict):
50
                for k2, v2 in v.items():
51
                    k3 = "_".join([k, k2])
52
                    tb_wandb_log(
53
                        f"eval/{k3}",
54
                        v2,
55
                        neox_args.iteration,
56
                        use_wandb=neox_args.use_wandb,
57
                    )
58
            else:
59
                tb_wandb_log(
60
                    f"eval/{k}",
61
                    v,
62
                    neox_args.iteration,
63
                    use_wandb=neox_args.use_wandb,
64
                )
65

66
        pprint(results)
67
        results_path = (
68
            f'eval_results_{datetime.now().strftime("%m-%d-%Y-%H-%M-%S")}.json'
69
        )
70
        if neox_args.eval_results_prefix:
71
            results_path = f"{neox_args.eval_results_prefix}_{results_path}"
72
        with open(results_path, "w") as f:
73
            json.dump(results, f, indent=4)
74

75

76
if __name__ == "__main__":
77
    main()
78