def METHOD_NAME(
    context: Context = None,
    swds: str = "mkqa/version/latest",
) -> None:
    checkpoint = str(ROOTDIR / "models")
    if not os.path.exists(checkpoint):
        from download_model import download
        download(os.getenv("BLOOM_MODEL_NAME") or "bigscience/bloom-560m")
    tokeniser = BloomTokenizerFast.from_pretrained(checkpoint, add_prefix_space=True)
    model = BloomForCausalLM.from_pretrained(
        checkpoint, torch_dtype="auto", device_map="auto"
    )
    swds_name = context.dataset_uris[0] if context else swds
    sw_dataset = dataset(swds_name, readonly=True, create="forbid")
    sw_dataset = sw_dataset.with_loader_config(
        field_transformer=ds_key_selectors.get(sw_dataset._uri.name, None)
    )
    hgds = swds2hgds(sw_dataset)
    input_ids = tokenise_data(hgds, tokeniser)
    model.gradient_checkpointing_enable()
    model.is_parallelizable = True
    model.model_parallel = True
    # train
    trainer = ModifiedTrainer(
        model=model,
        train_dataset=input_ids,
        args=TrainingArguments(
            output_dir=checkpoint,
            per_device_train_batch_size=int(os.getenv("BLOOM_TRAIN_BATCH_SIZE")) or 8,
            gradient_accumulation_steps=1,
            num_train_epochs=int(os.getenv("BLOOM_TRAIN_EPOCHES")) or 4,
            learning_rate=2e-5,
            fp16=True,
            logging_steps=10,
        ),
        data_collator=data_collator,
    )
    trainer.train()
    swmp.build(
        workdir=ROOTDIR,
        name=os.getenv("BLOOM_MODEL_NAME") or "bloom",
        modules=[ft],
    )