Move fqn mapping logic to StateDictAdapter

[wesleytruong]

This moves the logic that parses model.safetensors.index.json and generates the fqn_to_index_mapping to StateDictAdapter since this logic should be shared by all classes that inherit from StateDictAdapter.

[tianyu-l]

maybe don't need this BaseStateDictAdapter -- could you think of a case where people would like to inherit BaseStateDictAdapter but not StateDictAdapter? What do you think

[wesleytruong]

I'm not sure if it would be better to handle state dict adapter this way, but there are some multi-modal repositories that may have multiple safetensors.index.jsons such as even https://huggingface.co/black-forest-labs/FLUX.1-dev/tree/main. I think you could handle such cases by using sd_adapter on each of the individual sub-models, but perhaps someone may prefer having them all under the same class? In the case of Flux, the extra models are just autoencoders that are not being trained, but perhaps someone may want to use torchtitan in a way that trains multiple models at the same time

[tianyu-l]

In the case of FLUX, how would it work?

1. Does our approach still work? It sounds to me that we are not ready to read https://huggingface.co/black-forest-labs/FLUX.1-dev/tree/main/transformer
2. We don't support loading from multiple models from multiple folders anyway, so not sure if overgeneralizing helps. In the future if we support that, we probably need multiple StateDictAdapter?

All that said I'm OK with this change, but probably need to change https://github.com/pytorch/torchtitan/blob/main/torchtitan/protocols/train_spec.py#L24 to use the base one.

[wesleytruong]

I think our approach will still work currently for all models using tokenizer_path since its logic for downloading tokenizer files is identical to previous download_tokenizer script. For loading the model in Flux we can pass the hf_assets_path of the full repo to FluxStateDictAdapter and then FluxStateDictAdapter can pass hf_assets_path + "transformer" to the parent class. I agree though that I think I can iterate on the download_hf_assets options to narrow search to certain subfolders, or pattern match on full path names instead of base names.

[tianyu-l]

sounds good.

please resolve with @wwwjn on parallel work #1538