Update Gemma 7b tests to use optionally internal GCS buckets for testing

[A9isha]

For models such as Gemma-7B which needs multihost TPUs for running train or decode, we can run the checkpoint conversion step of creating the MaxText compatible Orbax checkpoint in a CPU machine or TPU v4-8 (basically any single host machine with enough RAM) in end_to_end/gemma/7b/1_test_gemma.sh, and run the decode and train steps on a multi-host TPU in end_to_end/gemma/7b/2_test_gemma.sh.

When connecting the runs of these two separate scripts:
Run the two scripts as follows:

export BASE_OUTPUT_PATH=/path/to/GCS/bucket; bash end_to_end/tpu/gemma/7b/1_test_gemma.sh
export BASE_OUTPUT_PATH=/path/to/same/GCS/bucket; bash end_to_end/tpu/gemma/7b/2_test_gemma.sh

[jonb377]

LGTM, thanks Anisha!

[jonb377]

Just confirming, idx has been absorbed into MODEL_BUCKET, so this will still be distinct for each run?

[A9isha]

Right, since MODEL_BUCKET has the idx absorbed into it, we are getting CONVERTED_CHECKPOINT distinct in each run. We are also pulling out idx in line 21 from MODEL_BUCKET, and used in the different RUN_NAME vars in line 24, 47, 56, these RUN_NAME vars would make the outputs of end_to_end/gemma/7b/2_test_gemma.sh distinct for each run.

Is this what you meant?

[rwitten]

what is this AWK doing? Seems unwieldy, maybe you can explain the problem (ideally in some message about the commit?)

[A9isha]

Thanks Rafi - I added the explanation in the description of this PR.

[rwitten]

Ok that is super helpful.

Is there a way we could not have thing branching? Like do we really to preserve multiple options here?

[A9isha]

Did you mean we remove the two options of connecting the two testing scripts, and only have the option of this?
(we need this option of connecting the two scripts for our internal testing though )

export MODEL_BUCKET=/path/to/GCS/bucket/and/an/index; bash end_to_end/gemma/7b/1_test_gemma.sh
export MODEL_BUCKET=/path/to/GCS/bucket/and/an/index; bash end_to_end/gemma/7b/2_test_gemma.sh

[rwitten]

Yes

[rwitten]

(Or maybe focus on a version that doesn't require parsing so you apss two args.

[A9isha]

I can definitely make passing MODEL_BUCKET as the only option

Could we pass both a GCS path as well as an additional index from airflow? @jonb377 and @RissyRan

[A9isha]

Followed up with @RissyRan - it is possible, it would need some modifications to our airflow+XLML tests code.

[A9isha]

@rwitten I think the much simpler and cleaner solution will be to have one GCS path shared between the two scripts and it would be the BASE_OUTPUT_DIRECTORY, and write everything there: for e.g., scanned ckpt, unscanned ckpt, and all outputs of train and decode (in separate subfolders). This shared GCS path will be unique since it will have the datetime inside it and will be generated from airflow.

I updated the code, please take a look.

cc: @RissyRan @jonb377 @khatwanimohit

[khatwanimohit]

LGTM with a nit

[khatwanimohit]

Nit:
In MaxText land, BASE_OUTPUT_DIRECTORY is usually just a gcs bucket, we attach the run_name to it and create a unique path for the run.
I think we should be more clear in the example usage here, saying that BASE_OUTPUT_DIRECTORY is the full path of user's unique run.

I think we should also consider renaming that variable to BASE_OUTPUT_PATH to avoid confusion.

[A9isha]

Thanks Mohit!

I added a comment in the scripts to clarify and also renamed to BASE_OUTPUT_PATH.